bash
python -m torch.distributed.launch --nproc_per_node=2 examples/pytorch/translation/run_translation.py \
--model_name_or_path t5-small --per_device_train_batch_size 1   \
--output_dir output_dir --overwrite_output_dir \
--do_train --max_train_samples 500 --num_train_epochs 1 \
--dataset_name wmt16 --dataset_config "ro-en" \
--source_lang en --target_lang ro \
--fp16 --sharded_ddp zero_dp_2
zero_dp_2 is an optimized version of the simple wrapper, while zero_dp_3 fully shards model weights,
gradients and optimizer states.
Both are compatible with adding cpu_offload to enable ZeRO-offload (activate it like this: --sharded_ddp "zero_dp_2 cpu_offload").
Notes:
This feature requires distributed training (so multiple GPUs).
It is not implemented for TPUs.
It works with --fp16 too, to make things even faster.
