diff --git a/train.py b/train.py index 0cee768..19ad54a 100644 --- a/train.py +++ b/train.py @@ -169,17 +169,18 @@ def train(config_path): eval_steps=config.get("eval_steps", 100), bf16=True, # gradient_checkpointing removed - FSDP activation_checkpointing handles it - fsdp="full_shard", + fsdp=True, fsdp_config={ "sharding_strategy": "SHARD_GRAD_OP", "cpu_offload": False, "activation_checkpointing": True, "limit_all_gathers": True, - "sync_module_states": True, + "sync_module_states": False, # Don't sync from CPU, FSDP will handle placement "backward_prefetch": "backward_pre", "forward_prefetch": "true", "auto_wrap_policy": "TRANSFORMER_BASED_WRAP", "transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer", + "mixed_precision": {"param_dtype": torch.bfloat16, "reduce_dtype": torch.float32, "buffer_dtype": torch.float32}, }, )