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

feat: use accelerate device_map for DISTRIBUTED model loading (Test 7 method)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 15:19:04 -04:00
parent 48e44cdf1c
commit 4348323116
1 changed files with 40 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
train.py
View File

@@ -33,100 +33,48 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model with multiple strategies # Load model with accelerate device_map (Test 7 method - DISTRIBUTED)
print(f"\n[INFO] Loading {config['base_model']}...") print(f"\n[INFO] Loading {config['base_model']}...")
errors = []
# ------------------------------------------------------------------ # Detect layer names dynamically
# Strategy 1: Load already-quantized 4-bit model (distributed across GPUs) from transformers import AutoConfig
# ------------------------------------------------------------------ print(" Detecting layer names...")
print("\n[1/4] Trying: 4-bit model AS-IS (distributed across GPUs)...") config_model = AutoConfig.from_pretrained(config["base_model"], trust_remote_code=True)
try: layer_names = [f"{config_model.model_type.title().replace('_', '')}DecoderLayer"]
model = AutoModelForCausalLM.from_pretrained( if 'moe' in config_model.model_type.lower():
config["base_model"], layer_names.append(f"{config_model.model_type.title().replace('_', '')}SparseMoeBlock")
device_map="auto", # Distribute layers across GPUs automatically print(f" Detected layers: {layer_names}")
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: 4-bit model distributed across GPUs (no FSDP)")
except Exception as e:
errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {e}")
# -------------------------------------------------------------- # Use accelerate to create device_map
# Strategy 2: 4-bit model with FSDP (load to GPU, FSDP shards) from accelerate import infer_auto_device_map
# -------------------------------------------------------------- print(" Creating device_map with accelerate...")
print("\n[2/4] Trying: 4-bit model with FSDP (load to GPU)...")
try:
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
device_map="auto", # Load to GPU first
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: 4-bit model loaded to GPU (FSDP will shard)")
except Exception as e:
errors.append(("QLoRA 4-bit FSDP", e))
print(f"✗ Failed: {e}")
# -------------------------------------------------------------- # First load to CPU to get the model structure
# Strategy 3: BF16 CPU (model too large for single GPU) model_temp = AutoModelForCausalLM.from_pretrained(
# -------------------------------------------------------------- config["base_model"],
print("\n[3/4] Trying: bf16 CPU...") device_map="cpu",
try: torch_dtype=torch.float16,
model = AutoModelForCausalLM.from_pretrained( trust_remote_code=True,
config["base_model"], low_cpu_mem_usage=True,
torch_dtype=torch.bfloat16, )
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
print("✓ Success: bf16 CPU")
except Exception as e:
errors.append(("bf16 CPU", e))
print(f"✗ Failed: {e}")
# ---------------------------------------------------------- # Create device_map
# Strategy 4: FP16 CPU (fallback) device_map = infer_auto_device_map(
# ---------------------------------------------------------- model_temp,
print("\n[4/4] Trying: fp16 CPU...") max_memory={i: "15GB" for i in range(torch.cuda.device_count())},
try: no_split_module_classes=layer_names,
model = AutoModelForCausalLM.from_pretrained( )
config["base_model"], print(f" Created device_map with {len(device_map)} entries")
torch_dtype=torch.float16,
device_map="cpu",
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: fp16 CPU")
except Exception as e:
errors.append(("fp16 CPU", e))
print(f"✗ Failed: {e}")
# ------------------------------------------------------ # Reload with device_map
# Strategy 5: 4-bit AS-IS (already quantized) model = AutoModelForCausalLM.from_pretrained(
# ---------------------------------------------------------- config["base_model"],
print("\n[5/5] Trying: 4-bit AS-IS...") device_map=device_map,
try: torch_dtype=torch.float16,
model = AutoModelForCausalLM.from_pretrained( trust_remote_code=True,
config["base_model"], low_cpu_mem_usage=True,
torch_dtype=torch.float16, )
device_map="cpu", print("✓ Success: Model distributed across GPUs via accelerate device_map")
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: 4-bit AS-IS")
except Exception as e:
errors.append(("4-bit AS-IS", e))
print(f"✗ Failed: {e}")
msg = "\n".join(
f"{name}: {err}" for name, err in errors
)
raise RuntimeError(
f"All loading strategies failed:\n\n{msg}"
)
# Add LoRA # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(