diff --git a/training/scripts/train.py b/training/scripts/train.py index 8205fa4..c77198c 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -59,76 +59,72 @@ 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}") - raise RuntimeError("All loading strategies failed!") + 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 from peft import prepare_model_for_kbit_training