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

fix: load model to single GPU (MoE requires all layers on same device)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 17:50:01 -04:00
parent 0456d55c8a
commit 92999b7783
1 changed files with 9 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
train.py
View File

@@ -33,48 +33,18 @@ def train(config_path):
print(f"Loading model: {config['base_model']}")
# Load model with accelerate device_map (Test 7 method - DISTRIBUTED)
print(f"\n[INFO] Loading {config['base_model']}...")
# Detect layer names dynamically
from transformers import AutoConfig
print(" Detecting layer names...")
config_model = AutoConfig.from_pretrained(config["base_model"], trust_remote_code=True)
layer_names = [f"{config_model.model_type.title().replace('_', '')}DecoderLayer"]
if 'moe' in config_model.model_type.lower():
layer_names.append(f"{config_model.model_type.title().replace('_', '')}SparseMoeBlock")
print(f" Detected layers: {layer_names}")
# Use accelerate to create device_map
from accelerate import infer_auto_device_map
print(" Creating device_map with accelerate...")
# First load to CPU to get the model structure
model_temp = AutoModelForCausalLM.from_pretrained(
config["base_model"],
device_map="cpu",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
# Create device_map
device_map = infer_auto_device_map(
model_temp,
max_memory={i: "15GB" for i in range(torch.cuda.device_count())},
no_split_module_classes=layer_names,
)
print(f" Created device_map with {len(device_map)} entries")
# Reload with device_map
# Load model to single GPU (MoE needs all layers on same device)
print(f"\n[INFO] Loading {config['base_model']} to single GPU...")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
device_map=device_map,
device_map="cuda:0",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: Model distributed across GPUs via accelerate device_map")
print("✓ Success: Model loaded to GPU 0")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
# Add LoRA
lora_config = LoraConfig(
@@ -107,12 +77,10 @@ def train(config_path):
},
)
# Model is already distributed across GPUs via device_map
# No FSDP needed - device_map handles distribution
print("✓ Model already distributed across GPUs (device_map)")
# Model is on single GPU
print("✓ Model loaded to single GPU")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
if torch.cuda.device_count() > 1:
print(f" GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
# Training arguments
training_args = TrainingArguments(