cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: FSDP loads model to GPU first, then shards across GPUs

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 12:17:44 -04:00
parent f9c748706f
commit 7f498134f2
1 changed files with 9 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
train.py
View File

@@ -55,25 +55,18 @@ def train(config_path):
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# -------------------------------------------------------------- # --------------------------------------------------------------
# Strategy 2: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs) # Strategy 2: 4-bit model with FSDP (load to GPU, FSDP shards)
# -------------------------------------------------------------- # --------------------------------------------------------------
print("\n[2/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...") print("\n[2/4] Trying: 4-bit model with FSDP (load to GPU)...")
try: try:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_storage=torch.bfloat16, # Enable FSDP sharding of 4-bit weights
)
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, device_map="auto", # Load to GPU first
device_map="cpu", # Load to CPU, FSDP shards later torch_dtype=torch.float16,
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 (FSDP will shard across GPUs)") print("✓ Success: 4-bit model loaded to GPU (FSDP will shard)")
except Exception as e: except Exception as e:
errors.append(("QLoRA 4-bit FSDP", e)) errors.append(("QLoRA 4-bit FSDP", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
@@ -188,14 +181,14 @@ def train(config_path):
transformer_layer_cls={Qwen3_5MoeDecoderLayer}, transformer_layer_cls={Qwen3_5MoeDecoderLayer},
) )
# Wrap model with FSDP on CPU first # Wrap model with FSDP on GPU
print("Wrapping model with FSDP on CPU...") print("Wrapping model with FSDP on GPU...")
model = FSDP( model = FSDP(
model, model,
auto_wrap_policy=get_auto_wrap_policy(model), auto_wrap_policy=get_auto_wrap_policy(model),
device_id=None, # Keep on CPU initially device_id=torch.cuda.current_device(), # Keep on current GPU
mixed_precision=None, mixed_precision=None,
sync_module_states=False, # Model is on CPU, no sync needed sync_module_states=False, # Model is already on GPU
use_orig_params=True, use_orig_params=True,
) )
print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)") print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")