diff --git a/train.py b/train.py index a1f8ee1..165c813 100644 --- a/train.py +++ b/train.py @@ -38,9 +38,9 @@ def train(config_path): errors = [] # ------------------------------------------------------------------ - # Strategy 1: QLoRA with FSDP (preferred) + # Strategy 1: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs) # ------------------------------------------------------------------ - print("\n[1/4] Trying: 4-bit QLoRA (FSDP)...") + print("\n[1/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...") try: bnb_config = BitsAndBytesConfig( load_in_4bit=True, @@ -56,7 +56,7 @@ def train(config_path): trust_remote_code=True, low_cpu_mem_usage=True, ) - print("✓ Success: QLoRA 4-bit loaded to CPU") + print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)") except Exception as e: errors.append(("QLoRA 4-bit", e)) print(f"✗ Failed: {e}") @@ -156,6 +156,12 @@ def train(config_path): 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 @@ -175,7 +181,7 @@ def train(config_path): sync_module_states=True, use_orig_params=True, ) - print("✓ Model wrapped with FSDP") + print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)") # Training arguments (no FSDP config - we handle it manually) training_args = TrainingArguments(