diff --git a/train-on-this-server.sh b/train-on-this-server.sh index a07f3a7..d1dc4ce 100755 --- a/train-on-this-server.sh +++ b/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!" \ No newline at end of file diff --git a/train.py b/train.py index a2f9370..10761ed 100644 --- a/train.py +++ b/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 - - # 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}, - ) - - # 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"], - 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 - ) - - # SFT Trainer - from trl import SFTTrainer - - 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 + # 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") + + # 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), + ) + + # SFT Trainer + from trl import SFTTrainer + + trainer = SFTTrainer( + model=model, + processing_class=tokenizer, + train_dataset=dataset["train"], + eval_dataset=dataset.get("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']}") def main():