# Distibuted Data Parallel in SageMaker

**URL:** <https://discuss.huggingface.co/t/distibuted-data-parallel-in-sagemaker/71965>\
**Category:** Amazon SageMaker\
**Created:** [February 5, 2024, 6:02pm UTC](https://discuss.huggingface.co/t/distibuted-data-parallel-in-sagemaker/71965 "2024-02-05T18:02:41Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![ebruk](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/ebruk/32/27504_2.png) [@ebruk](https://discuss.huggingface.co/u/ebruk)\
**Post date:** [February 5, 2024, 6:02pm UTC](https://discuss.huggingface.co/t/distibuted-data-parallel-in-sagemaker/71965/1 "2024-02-05T18:02:41Z")

</div>

I have this error when try to run the custom dataset with distributed train: ErrorMessage "ModuleNotFoundError: No module named ‘torch.distributed.\_functional\_collectives’

I couldn’t find what caused this error.

The estimator :  
distribution ={“pytorchddp”: {“enabled”: True}}

# hyperparameters, which are passed into the training job

hyperparameters ={  
#‘model\_id’: model\_id, # pre-trained model  
‘dataset\_path’: dataset\_path, # path where sagemaker will save training dataset  
‘epochs’: epochs, # number of training epochs  
‘per\_device\_train\_batch\_size’: per\_device\_train\_batch\_size, # batch size for training  
‘lr’: learning\_rate, # learning rate used during training  
}

estimator = PyTorch(  
checkpoint\_s3\_uri = f’s3://{sess.default\_bucket()}/v02/checkpoints\_dist’,  
entry\_point = ‘pylight\_sc.py’, # train script  
source\_dir = ‘scripts’, # directory which includes all the files needed for training  
instance\_type = ‘ml.p4d.24xlarge’, # instances type used for the training job  
instance\_count = 1, # the number of instances used for training  
max\_run = 2_24_60\*60, # the max time of train  
base\_job\_name = job\_name, # the name of the training job  
role = role, # Iam role used in training job to access AWS ressources, e.g. S3  
volume\_size = 300, # the size of the EBS volume in GB  
py\_version = “py310”,  
pytorch\_version = ‘2.0.1’, # the pytorch\_version version used in the training job  
framework\_version = ‘2.0.0’,  
hyperparameters = hyperparameters,  
distribution = distribution  
)

And the pylight\_sc.py:

def training\_function(args):  
# set seed  
set\_seed(args.seed)

```
env = LightningEnvironment()
env.world_size = lambda: int(os.environ["WORLD_SIZE"])
env.global_rank = lambda: int(os.environ["RANK"])

policy = partial(
    size_based_auto_wrap_policy, 
    min_num_params=10000
)

fsdp = FSDPStrategy(
    auto_wrap_policy=policy,
    process_group_backend="smddp", 
    cluster_environment=env
)

dataset = load_from_disk(args.dataset_path)
# load model from the hub with a bnb config
bnb_config = BitsAndBytesConfig(
    # load_in_8bit=True,
    load_in_4bit=True,
    # bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    args.model_id,
    use_cache=False if args.gradient_checkpointing else True, # this is needed for gradient checkpointing
    #trust_remote_code=True, # ATTENTION: This allows remote code execution
    device_map="auto",
    quantization_config=bnb_config
)

# create peft config
model = create_peft_config(model, args.gradient_checkpointing)
num_gpus = torch.cuda.device_count()

# Define training args
checkpoint_dir = "/opt/ml/checkpoints"
training_args = TrainingArguments(
    output_dir=checkpoint_dir,
    overwrite_output_dir=True,
    per_device_train_batch_size=args.per_device_train_batch_size,
    bf16=args.bf16, # Use BF16 if available
    save_total_limit=2, 
    learning_rate=args.lr,
    num_train_epochs=args.epochs,
    gradient_checkpointing=args.gradient_checkpointing,
    remove_unused_columns=False,
    # logging strategies
    logging_dir=f"{checkpoint_dir}/logs",
    logging_strategy="steps",
    logging_steps=10,
    save_strategy="steps",
    save_steps= 300,
    report_to=["wandb"],
)

# Create Trainer instance 
trainer = pl.Trainer(
    model=model,
    strategy=fsdp,
    args=training_args,
    devices=num_gpus,
    train_dataset=dataset,
    data_collator=default_data_collator,
)

# Start training
trainer.train(resume_from_checkpoint= True)

```