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
-echo "Starting training with accelerate..."
+# Run training with single process (model already distributed via device_map)
+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,105 +96,63 @@ def train(config_path):
     from datasets import load_dataset
     import os
     
-    # Get dataset paths - training/data/ is relative to project root
-    repo_root = Path(__file__).parent
-    train_path = str(repo_root / "training" / "data" / "train.jsonl")
-    test_path = str(repo_root / "training" / "data" / "test.jsonl")
-    print(f"Looking for dataset at: {train_path}")
+    # Get dataset path from config
+    dataset_path = config["dataset"][0]["path"]
+    print(f"Loading dataset from: {dataset_path}")
     
     dataset = load_dataset(
         "json",
         data_files={
-            "train": train_path,
-            "test": test_path,
+            "train": dataset_path,
         },
     )
     
-    # Manually wrap model with FSDP before passing to trainer
-    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
+    # Model is already distributed across GPUs via device_map
+    # No FSDP needed - device_map handles distribution
+    print("✓ Model already distributed across GPUs (device_map)")
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    if torch.cuda.device_count() > 1:
+        print(f"   GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
     
-        # Initialize distributed process group (required for FSDP)
-        if not dist.is_initialized():
-            dist.init_process_group(backend="nccl")
-            print("✓ Distributed process group initialized")
+    # Training arguments
+    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=config.get("gradient_checkpointing", True),
+    )
     
-        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},
-            )
+    # SFT Trainer
+    from trl import SFTTrainer
     
-        # 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)")
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset.get("test"),
+        args=training_args,
+    )
     
-        # 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
-        )
+    # Train
+    print("Starting training...")
+    trainer.train()
     
-        # SFT Trainer
-        from trl import SFTTrainer
+    # Save
+    trainer.save_model(config["train_params"]["output_dir"])
+    tokenizer.save_pretrained(config["train_params"]["output_dir"])
     
-        trainer = SFTTrainer(
-            model=model,
-            processing_class=tokenizer,
-            train_dataset=dataset["train"],
-            eval_dataset=dataset["test"],
-            args=training_args,
-        )
-        
-        # Train
-        print("Starting training...")
-        trainer.train()
-        
-        # Save
-        trainer.save_model(config["train_params"]["output_dir"])
-        tokenizer.save_pretrained(config["train_params"]["output_dir"])
-        
-        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
+    print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
 
 
 def main():