fix: use SHARD_GRAD_OP instead of FULL_SHARD for 2-GPU setup

2026-07-02 00:58:41 -04:00
parent e3ea45fe96
commit dbc97138e9
1 changed files with 4 additions and 2 deletions
--- a/train.py
+++ b/train.py
@@ -169,12 +169,14 @@ def train(config_path):
        # gradient_checkpointing removed - FSDP activation_checkpointing handles it
        fsdp=True,
        fsdp_config={
+            "sharding_strategy": "SHARD_GRAD_OP",
+            "cpu_offload": False,
+            "activation_checkpointing": True,
+            "limit_all_gathers": True,
            "backward_prefetch": "backward_pre",
            "forward_prefetch": "true",
-            "cpu_offload": "false",
            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
            "transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
-            "activation_checkpointing": True,
        },
    )