From 313b44381ffe93b2cf51284fab9b2dc06ec450a0 Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:10:54 -0400 Subject: [PATCH] fix: remove DeepSpeed, use torchrun without distributed training --- training/scripts/train.py | 67 ++++++++------------------------------- 1 file changed, 14 insertions(+), 53 deletions(-) diff --git a/training/scripts/train.py b/training/scripts/train.py index c77198c..947c61b 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -59,72 +59,33 @@ def train(config_path): except Exception as e: print(f"✗ Failed: {e}") - # Strategy 1: 4-bit model variants - print("\n[1/6] Trying: 4-bit model AS-IS with DeepSpeed...") + # Strategy 1: Load 4-bit model AS-IS (no quantization, no DeepSpeed) + print("\n[1/2] Trying: 4-bit model AS-IS...") try: - import deepspeed model = AutoModelForCausalLM.from_pretrained( config["base_model"], torch_dtype=torch.float16, device_map="auto", 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 model AS-IS") + print("✓ Success: 4-bit model loaded") except Exception as e: print(f"✗ Failed: {e}") - print("\n[2/6] Trying: 4-bit model with bf16 compute...") + # Strategy 2: bf16 to CPU + print("\n[2/2] Trying: bf16 model to CPU...") try: - 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: 4-bit with bf16 compute") + model = AutoModelForCausalLM.from_pretrained( + config["base_model"], + torch_dtype=torch.bfloat16, + device_map="cpu", + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + print("✓ Success: bf16 model loaded to CPU") 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!") + raise RuntimeError("All loading strategies failed!") # Prepare model for k-bit training from peft import prepare_model_for_kbit_training