fix: clarify strategy naming + init distributed process group for FSDP

This commit is contained in:
Christian Medina
2026-07-02 11:30:34 -04:00
parent 8cabc0e986
commit 93ac7391dc

View File

@@ -38,9 +38,9 @@ def train(config_path):
errors = [] 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: try:
bnb_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
@@ -56,7 +56,7 @@ def train(config_path):
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=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: except Exception as e:
errors.append(("QLoRA 4-bit", e)) errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {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 import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial 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): def get_auto_wrap_policy(model):
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer 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, sync_module_states=True,
use_orig_params=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 arguments (no FSDP config - we handle it manually)
training_args = TrainingArguments( training_args = TrainingArguments(