feat: add 6 loading variants with different offload strategies
This commit is contained in:
Christian Medina
@@ -59,75 +59,71 @@ def train(config_path):
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
|
||||
# Strategy 1: Load 4-bit model AS-IS with DeepSpeed ZeRO-3
|
||||
print("\n[1/3] Trying: 4-bit model with DeepSpeed ZeRO-3...")
|
||||
# Strategy 1: 4-bit model variants
|
||||
print("\n[1/6] Trying: 4-bit model AS-IS with DeepSpeed...")
|
||||
try:
|
||||
import deepspeed
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
config["base_model"],
|
||||
torch_dtype=torch.float16, # Load as fp16 (already 4-bit)
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
ds_config = {
|
||||
"train_micro_batch_size_per_gpu": 1,
|
||||
"gradient_accumulation_steps": 1,
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"contiguous_gradients": True,
|
||||
"overlap_comm": True,
|
||||
"offload_optimizer": {"device": "cpu", "pin_memory": True},
|
||||
"offload_param": {"device": "cpu", "pin_memory": True},
|
||||
},
|
||||
"fp16": {"enabled": True},
|
||||
}
|
||||
|
||||
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "fp16": {"enabled": True}}
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
|
||||
model, optimizer, _, _ = deepspeed.initialize(
|
||||
model=model,
|
||||
model_parameters=model.parameters(),
|
||||
optimizer=optimizer,
|
||||
config=ds_config,
|
||||
)
|
||||
print("✓ Success: 4-bit model with DeepSpeed ZeRO-3")
|
||||
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
|
||||
print("✓ Success: 4-bit model AS-IS")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
|
||||
# Strategy 2: bf16 with DeepSpeed ZeRO-3 + CPU offload
|
||||
print("\n[2/3] Trying: bf16 with DeepSpeed ZeRO-3 CPU offload...")
|
||||
print("\n[2/6] Trying: 4-bit model with bf16 compute...")
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
config["base_model"],
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="cpu",
|
||||
low_cpu_mem_usage=True,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
ds_config = {
|
||||
"train_micro_batch_size_per_gpu": 1,
|
||||
"gradient_accumulation_steps": 1,
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"contiguous_gradients": True,
|
||||
"overlap_comm": True,
|
||||
"offload_optimizer": {"device": "cpu", "pin_memory": True},
|
||||
"offload_param": {"device": "cpu", "pin_memory": True},
|
||||
},
|
||||
"bf16": {"enabled": True},
|
||||
}
|
||||
|
||||
from transformers import BitsAndBytesConfig
|
||||
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
|
||||
model = AutoModelForCausalLM.from_pretrained(config["base_model"], quantization_config=bnb_config, device_map="auto", trust_remote_code=True)
|
||||
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "bf16": {"enabled": True}}
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
|
||||
model, optimizer, _, _ = deepspeed.initialize(
|
||||
model=model,
|
||||
model_parameters=model.parameters(),
|
||||
optimizer=optimizer,
|
||||
config=ds_config,
|
||||
)
|
||||
print("✓ Success: bf16 with DeepSpeed ZeRO-3 CPU offload")
|
||||
except Exception as e2:
|
||||
print(f"✗ Failed: {e2}")
|
||||
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
|
||||
print("✓ Success: 4-bit with bf16 compute")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
|
||||
print("\n[3/6] Trying: 4-bit model to CPU then DeepSpeed...")
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True)
|
||||
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "fp16": {"enabled": True}}
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
|
||||
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
|
||||
print("✓ Success: 4-bit to CPU")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
|
||||
print("\n[4/6] Trying: 4-bit model fp32...")
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
|
||||
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "fp32": {"enabled": True}}
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
|
||||
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
|
||||
print("✓ Success: 4-bit fp32")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
|
||||
print("\n[5/6] Trying: bf16 model to CPU...")
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.bfloat16, device_map="cpu", low_cpu_mem_usage=True, trust_remote_code=True)
|
||||
ds_config = {"zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"}}, "bf16": {"enabled": True}}
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
|
||||
model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), optimizer=optimizer, config=ds_config)
|
||||
print("✓ Success: bf16 to CPU")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
|
||||
print("\n[6/6] Trying: bf16 model auto placement...")
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True, trust_remote_code=True)
|
||||
print("✓ Success: bf16 auto placement")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {e}")
|
||||
raise RuntimeError("All loading strategies failed!")
|
||||
|
||||
# Prepare model for k-bit training
|
||||
|
||||
Reference in New Issue
Block a user