feat: FSDP training failures now fallback to next strategy

train.py

@@ -176,76 +176,90 @@ def train(config_path):
     )
     
     # Manually wrap model with FSDP before passing to trainer
-    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
-    from functools import partial
-    import torch.distributed as dist
+    try:
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+        from functools import partial
+        import torch.distributed as dist
         
-    # Initialize distributed process group (required for FSDP)
-    if not dist.is_initialized():
-        dist.init_process_group(backend="nccl")
-        print("✓ Distributed process group initialized")
+        # Initialize distributed process group (required for FSDP)
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+            print("✓ Distributed process group initialized")
         
-    def get_auto_wrap_policy(model):
-        from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
-        return partial(
-            transformer_auto_wrap_policy,
-            transformer_layer_cls={Qwen3_5MoeDecoderLayer},
-            min_num_params=1,
+        def get_auto_wrap_policy(model):
+            from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
+            return partial(
+                transformer_auto_wrap_policy,
+                transformer_layer_cls={Qwen3_5MoeDecoderLayer},
+            )
+        
+        # Wrap model with FSDP on CPU first
+        print("Wrapping model with FSDP on CPU...")
+        model = FSDP(
+            model,
+            auto_wrap_policy=get_auto_wrap_policy(model),
+            device_id=None,  # Keep on CPU initially
+            mixed_precision=None,
+            sync_module_states=True,
+            use_orig_params=True,
+        )
+        print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")
+        
+        # Training arguments (no FSDP config - we handle it manually)
+        training_args = TrainingArguments(
+            output_dir=config["train_params"]["output_dir"],
+            num_train_epochs=config["train_params"]["num_train_epochs"],
+            per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
+            gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
+            learning_rate=float(config["train_params"]["learning_rate"]),
+            lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
+            weight_decay=config["train_params"]["weight_decay"],
+            warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]),
+            logging_steps=config["train_params"]["logging_steps"],
+            save_steps=config["train_params"]["save_steps"],
+            save_total_limit=config["train_params"]["save_total_limit"],
+            eval_strategy=config.get("eval_strategy", "steps"),
+            eval_steps=config.get("eval_steps", 100),
+            bf16=True,
+            gradient_checkpointing=False,  # FSDP handles it
+            # No fsdp config - we wrap manually above
         )
         
-    # Wrap model with FSDP on CPU first
-    print("Wrapping model with FSDP on CPU...")
-    model = FSDP(
-        model,
-        auto_wrap_policy=get_auto_wrap_policy(model),
-        device_id=None,  # Keep on CPU initially
-        mixed_precision=None,
-        sync_module_states=True,
-        use_orig_params=True,
-    )
-    print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")
+        # SFT Trainer
+        from trl import SFTTrainer
         
-    # Training arguments (no FSDP config - we handle it manually)
-    training_args = TrainingArguments(
-        output_dir=config["train_params"]["output_dir"],
-        num_train_epochs=config["train_params"]["num_train_epochs"],
-        per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
-        gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
-        learning_rate=float(config["train_params"]["learning_rate"]),
-        lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
-        weight_decay=config["train_params"]["weight_decay"],
-        warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]),
-        logging_steps=config["train_params"]["logging_steps"],
-        save_steps=config["train_params"]["save_steps"],
-        save_total_limit=config["train_params"]["save_total_limit"],
-        eval_strategy=config.get("eval_strategy", "steps"),
-        eval_steps=config.get("eval_steps", 100),
-        bf16=True,
-        gradient_checkpointing=False,  # FSDP handles it
-        # No fsdp config - we wrap manually above
-    )
+        trainer = SFTTrainer(
+            model=model,
+            processing_class=tokenizer,
+            train_dataset=dataset["train"],
+            eval_dataset=dataset["test"],
+            args=training_args,
+        )
         
-    # SFT Trainer
-    from trl import SFTTrainer
+        # Train
+        print("Starting training...")
+        trainer.train()
         
-    trainer = SFTTrainer(
-        model=model,
-        processing_class=tokenizer,
-        train_dataset=dataset["train"],
-        eval_dataset=dataset["test"],
-        args=training_args,
-    )
+        # Save
+        trainer.save_model(config["train_params"]["output_dir"])
+        tokenizer.save_pretrained(config["train_params"]["output_dir"])
         
-    # Train
-    print("Starting training...")
-    trainer.train()
+        print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
         
-    # Save
-    trainer.save_model(config["train_params"]["output_dir"])
-    tokenizer.save_pretrained(config["train_params"]["output_dir"])
+        # Cleanup distributed process group
+        if dist.is_initialized():
+            dist.destroy_process_group()
         
-    print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
+        # Return early - success
+        return
+    except Exception as e:
+        errors.append(("QLoRA 4-bit FSDP training", e))
+        print(f"✗ Failed during FSDP training: {e}")
+        # Cleanup distributed process group if initialized
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        # Fall through to next strategy
 
 
 def main():