batch was 32 on 1 V100 so it’s ok to scale 4x on 4*V100 right? my understanding is that tran_batch_size is the cluster-level batch, not the per-GPU batch right? here is my config. it ran fine both on p3.16x and p4d.24xl, and GPU mem is not 800%
from sagemaker.huggingface import HuggingFace

# hyperparameters, which are passed into the training job
hyperparameters={'epochs': 3,
                 'train_batch_size': 128,
                 'model_name':'distilbert-base-uncased',
                 'fp16': True}

huggingface_estimator = HuggingFace(entry_point='train.py',
                            source_dir='./scripts',
                            instance_type='ml.p4d.24xlarge',
                            instance_count=1,
                            role=role,
                            transformers_version='4.6',
                            pytorch_version='1.7',
                            py_version='py36',
                            hyperparameters = hyperparameters)

train_batch_size = 256 also works fine, it is as fast as 128 on p3.16x and 7% faster than 128 on p4d.24xlarge (both still fp16)