| It is recommended to set these values to auto and allow the [Trainer] to automatically assign the values. | |
| yml | |
| { | |
| "zero_optimization": { | |
| "stage": 3, | |
| "offload_optimizer": { | |
| "device": "cpu", | |
| "pin_memory": true | |
| }, | |
| "offload_param": { | |
| "device": "cpu", | |
| "pin_memory": true | |
| }, | |
| "overlap_comm": true, | |
| "contiguous_gradients": true, | |
| "sub_group_size": 1e9, | |
| "reduce_bucket_size": "auto", | |
| "stage3_prefetch_bucket_size": "auto", | |
| "stage3_param_persistence_threshold": "auto", | |
| "stage3_max_live_parameters": 1e9, | |
| "stage3_max_reuse_distance": 1e9, | |
| "stage3_gather_16bit_weights_on_model_save": true | |
| } | |
| } | |
| You can use the deepspeed.zero.Init context manager to initialize a model faster: | |
| from transformers import T5ForConditionalGeneration, T5Config | |
| import deepspeed | |
| with deepspeed.zero.Init(): | |
| config = T5Config.from_pretrained("google-t5/t5-small") | |
| model = T5ForConditionalGeneration(config) | |
| For pretrained models, the DeepSped config file needs to have is_deepspeed_zero3_enabled: true setup in [TrainingArguments] and it needs a ZeRO configuration enabled. |