diff --git a/train.py b/train.py index 31e2bb5..e3fdbed 100644 --- a/train.py +++ b/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 - - # 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, + 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 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, ) - - # 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 - ) - - # 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']}") + 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 def main():