fix: single-process training, remove FSDP (model pre-distributed via device_map)

train-on-this-server.sh

@@ -18,9 +18,9 @@ pip install --upgrade pip
 echo "Installing training dependencies..."
 pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
 
-# Run training with accelerate data parallelism
+# Run training with single process (model already distributed via device_map)
-echo "Starting training with accelerate..."
+echo "Starting training (single process, model pre-distributed)..."
 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
-accelerate launch --multi_gpu --num_processes 2 --mixed_precision bf16 train.py --config training/configs/ornith-35b-lora.yaml
+python train.py --config training/configs/ornith-35b-lora.yaml
 
 echo "Training completed!"

train.py

@@ -96,52 +96,25 @@ def train(config_path):
     from datasets import load_dataset
     import os
     
-    # Get dataset paths - training/data/ is relative to project root
+    # Get dataset path from config
-    repo_root = Path(__file__).parent
+    dataset_path = config["dataset"][0]["path"]
-    train_path = str(repo_root / "training" / "data" / "train.jsonl")
+    print(f"Loading dataset from: {dataset_path}")
-    test_path = str(repo_root / "training" / "data" / "test.jsonl")
-    print(f"Looking for dataset at: {train_path}")
     
     dataset = load_dataset(
         "json",
         data_files={
-            "train": train_path,
+            "train": dataset_path,
-            "test": test_path,
         },
     )
     
-    # Manually wrap model with FSDP before passing to trainer
+    # Model is already distributed across GPUs via device_map
-    try:
+    # No FSDP needed - device_map handles distribution
-        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    print("✓ Model already distributed across GPUs (device_map)")
-        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
-        from functools import partial
+    if torch.cuda.device_count() > 1:
-        import torch.distributed as dist
+        print(f"   GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
     
-        # Initialize distributed process group (required for FSDP)
+    # Training arguments
-        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},
-            )
-        
-        # Wrap model with FSDP on GPU
-        print("Wrapping model with FSDP on GPU...")
-        model = FSDP(
-            model,
-            auto_wrap_policy=get_auto_wrap_policy(model),
-            device_id=torch.cuda.current_device(),  # Keep on current GPU
-            mixed_precision=None,
-            sync_module_states=False,  # Model is already on GPU
-            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"],
@@ -157,8 +130,7 @@ def train(config_path):
         eval_strategy=config.get("eval_strategy", "steps"),
         eval_steps=config.get("eval_steps", 100),
         bf16=True,
-            gradient_checkpointing=False,  # FSDP handles it
+        gradient_checkpointing=config.get("gradient_checkpointing", True),
-            # No fsdp config - we wrap manually above
     )
     
     # SFT Trainer
@@ -168,7 +140,7 @@ def train(config_path):
         model=model,
         processing_class=tokenizer,
         train_dataset=dataset["train"],
-            eval_dataset=dataset["test"],
+        eval_dataset=dataset.get("test"),
         args=training_args,
     )
     
@@ -182,20 +154,6 @@ def train(config_path):
     
     print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
 
-        # Cleanup distributed process group
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        
-        # 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():
     parser = argparse.ArgumentParser(description="Train LoRA adapter")