From dfe85c3c82d1d47245b5a922190d9fef3cfe2e41 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 00:16:33 -0400
Subject: [PATCH] fix: fsdp=True instead of list + add activation_checkpointing

---
 train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index a0d0a61..6d348ee 100644
--- a/train.py
+++ b/train.py
@@ -167,13 +167,14 @@ def train(config_path):
         eval_steps=config.get("eval_steps", 100),
         bf16=True,
         gradient_checkpointing=config.get("gradient_checkpointing", True),
-        fsdp=["full_shard", "auto_wrap"],
+        fsdp=True,
         fsdp_config={
             "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,
         },
     )