## Below is an example yaml for BF16 mixed-precision training using PyTorch Fully Sharded Data Parallism (FSDP) with CPU offloading on 8 GPUs.
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
+distributed_type: FSDP
downcast_bf16: 'no'
dynamo_backend: 'NO'
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: T5Block
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
mixed_precision: bf16
num_machines: 1
+num_processes: 8
rdzv_backend: static
same_network: true
use_cpu: false
##
from accelerate import Accelerator

def main():
    accelerator = Accelerator()
-    model, optimizer, dataloader, scheduler = accelerator.prepare(
-      model, optimizer, dataloader, scheduler
-    )
+    model = accelerator.prepare(model)
+    # Optimizer can be any PyTorch optimizer class
+    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+    optimizer, dataloader, scheduler = accelerator.prepare(
+      optimizer, dataloader, scheduler
+    )
    
    ...

    accelerator.unwrap_model(model).save_pretrained(
            args.output_dir,
            is_main_process=accelerator.is_main_process,
            save_function=accelerator.save,
+            state_dict=accelerator.get_state_dict(model)
    )
    ...
## If the YAML was generated through the `accelerate config` command: ``` accelerate launch {script_name.py} {--arg1} {--arg2} ... ``` If the YAML is saved to a `~/config.yaml` file: ``` accelerate launch --config_file ~/config.yaml {script_name.py} {--arg1} {--arg2} ... ``` Or you can use `accelerate launch` with right configuration parameters and have no `config.yaml` file: ``` accelerate launch \ --use_fsdp \ --num_processes=8 \ --mixed_precision=bf16 \ --fsdp_sharding_strategy=1 \ --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP \ --fsdp_transformer_layer_cls_to_wrap=T5Block \ --fsdp_offload_params=true \ {script_name.py} {--arg1} {--arg2} ... ``` ## For PyTorch FDSP, you need to prepare the model first **before** preparing the optimizer since FSDP will shard parameters in-place and this will break any previously initialized optimizers. For transformer models, please use `TRANSFORMER_BASED_WRAP` auto wrap policy as shown in the config above. ## To learn more checkout the related documentation: - How to use Fully Sharded Data Parallelism - Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel