fix: disable sync_module_states + add mixed precision for FSDP

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},
         },
     )