refactor: comment out failing tests, add Test 7-10 variations
This commit is contained in:
Christian Medina
@@ -30,194 +30,194 @@ def check_gpu_memory():
|
||||
else:
|
||||
return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
|
||||
|
||||
def test_strategy_1():
|
||||
"""Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
|
||||
print("=" * 80)
|
||||
# def test_strategy_1():
|
||||
# """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
|
||||
# print("\n" + "=" * 80)
|
||||
# print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
|
||||
# print("=" * 80)
|
||||
#
|
||||
# try:
|
||||
# print(" Loading bf16 model with BnB 4-bit...")
|
||||
# bnb_config = BitsAndBytesConfig(
|
||||
# load_in_4bit=True,
|
||||
# bnb_4bit_quant_type="nf4",
|
||||
# bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
# )
|
||||
# model = AutoModelForCausalLM.from_pretrained(
|
||||
# "/data/models/Ornith-1.0-35B", # ← bf16 model
|
||||
# quantization_config=bnb_config,
|
||||
# device_map="auto",
|
||||
# trust_remote_code=True,
|
||||
# low_cpu_mem_usage=True,
|
||||
# )
|
||||
# print(" ✓ Model loaded successfully")
|
||||
#
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f"\n Pattern: {pattern}")
|
||||
# return True, pattern
|
||||
# except Exception as e:
|
||||
# print(f"\n ✗ FAILED: {e}")
|
||||
# return False, str(e)
|
||||
|
||||
try:
|
||||
print(" Loading bf16 model with BnB 4-bit...")
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B", # ← bf16 model
|
||||
quantization_config=bnb_config,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded successfully")
|
||||
# def test_strategy_2():
|
||||
# """Test 2: bf16 model + BnB 4-bit (alternative config)"""
|
||||
# print("\n" + "=" * 80)
|
||||
# print("TEST 2: bf16 model + BnB 4-bit (alt config)")
|
||||
# print("=" * 80)
|
||||
#
|
||||
# try:
|
||||
# torch.cuda.empty_cache()
|
||||
# print(" Loading bf16 model with BnB 4-bit...")
|
||||
# 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(
|
||||
# "/data/models/Ornith-1.0-35B", # ← bf16 model
|
||||
# quantization_config=bnb_config,
|
||||
# device_map="auto",
|
||||
# trust_remote_code=True,
|
||||
# low_cpu_mem_usage=True,
|
||||
# )
|
||||
# print(" ✓ Model loaded successfully")
|
||||
#
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f"\n Pattern: {pattern}")
|
||||
# return True, pattern
|
||||
# except Exception as e:
|
||||
# print(f"\n ✗ FAILED: {e}")
|
||||
# return False, str(e)
|
||||
|
||||
pattern = check_gpu_memory()
|
||||
print(f"\n Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
return False, str(e)
|
||||
# def test_strategy_3():
|
||||
# """Test 3: device_map with explicit GPU assignment"""
|
||||
# print("\n" + "=" * 80)
|
||||
# print("TEST 3: device_map with explicit GPU assignment")
|
||||
# print("=" * 80)
|
||||
#
|
||||
# try:
|
||||
# torch.cuda.empty_cache()
|
||||
# print(" Loading model with explicit device_map...")
|
||||
#
|
||||
# # Get model config to determine layers
|
||||
# from transformers import AutoConfig
|
||||
# config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
|
||||
# num_layers = config.num_hidden_layers
|
||||
#
|
||||
# # Split layers: first half on GPU 0, second half on GPU 1
|
||||
# device_map = {}
|
||||
# for i in range(num_layers):
|
||||
# if i < num_layers // 2:
|
||||
# device_map[f"model.layers.{i}"] = 0
|
||||
# else:
|
||||
# device_map[f"model.layers.{i}"] = 1
|
||||
#
|
||||
# # Embeddings and norm on GPU 0
|
||||
# device_map["model.embed_tokens"] = 0
|
||||
# device_map["model.norm"] = 0
|
||||
# device_map["lm_head"] = 0
|
||||
#
|
||||
# print(f" Created device_map with {len(device_map)} entries")
|
||||
# model = AutoModelForCausalLM.from_pretrained(
|
||||
# "/data/models/Ornith-1.0-35B",
|
||||
# device_map=device_map,
|
||||
# torch_dtype=torch.bfloat16,
|
||||
# trust_remote_code=True,
|
||||
# low_cpu_mem_usage=True,
|
||||
# )
|
||||
# print(" ✓ Model loaded successfully")
|
||||
#
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f"\n Pattern: {pattern}")
|
||||
# return True, pattern
|
||||
# except Exception as e:
|
||||
# print(f"\n ✗ FAILED: {e}")
|
||||
# return False, str(e)
|
||||
|
||||
def test_strategy_2():
|
||||
"""Test 2: bf16 model + BnB 4-bit (alternative config)"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 2: bf16 model + BnB 4-bit (alt config)")
|
||||
print("=" * 80)
|
||||
# def test_strategy_4():
|
||||
# """Test 4: Load to CPU, then move to GPU manually"""
|
||||
# print("\n" + "=" * 80)
|
||||
# print("TEST 4: Load to CPU, then move to GPU")
|
||||
# print("=" * 80)
|
||||
#
|
||||
# try:
|
||||
# torch.cuda.empty_cache()
|
||||
# print(" Loading bf16 model to CPU...")
|
||||
# model = AutoModelForCausalLM.from_pretrained(
|
||||
# "/data/models/Ornith-1.0-35B",
|
||||
# device_map="cpu",
|
||||
# torch_dtype=torch.bfloat16,
|
||||
# trust_remote_code=True,
|
||||
# low_cpu_mem_usage=True,
|
||||
# )
|
||||
# print(" ✓ Model loaded to CPU")
|
||||
#
|
||||
# # Count params on CPU
|
||||
# cpu_params = sum(p.numel() for p in model.parameters() if p.device.type == 'cpu')
|
||||
# print(f" CPU parameters: {cpu_params / 1e9:.2f}B")
|
||||
#
|
||||
# # Move to GPU 0
|
||||
# print("\n Moving to GPU 0...")
|
||||
# model = model.to("cuda:0")
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f" Pattern after move to GPU 0: {pattern}")
|
||||
#
|
||||
# # Move to GPU 1
|
||||
# print("\n Moving to GPU 1...")
|
||||
# model = model.to("cuda:1")
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f" Pattern after move to GPU 1: {pattern}")
|
||||
#
|
||||
# return True, "LOADED_TO_CPU_THEN_GPU"
|
||||
# except Exception as e:
|
||||
# print(f"\n ✗ FAILED: {e}")
|
||||
# return False, str(e)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Loading bf16 model with BnB 4-bit...")
|
||||
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(
|
||||
"/data/models/Ornith-1.0-35B", # ← bf16 model
|
||||
quantization_config=bnb_config,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded successfully")
|
||||
|
||||
pattern = check_gpu_memory()
|
||||
print(f"\n Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_3():
|
||||
"""Test 3: device_map with explicit GPU assignment"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 3: device_map with explicit GPU assignment")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Loading model with explicit device_map...")
|
||||
|
||||
# Get model config to determine layers
|
||||
from transformers import AutoConfig
|
||||
config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
|
||||
num_layers = config.num_hidden_layers
|
||||
|
||||
# Split layers: first half on GPU 0, second half on GPU 1
|
||||
device_map = {}
|
||||
for i in range(num_layers):
|
||||
if i < num_layers // 2:
|
||||
device_map[f"model.layers.{i}"] = 0
|
||||
else:
|
||||
device_map[f"model.layers.{i}"] = 1
|
||||
|
||||
# Embeddings and norm on GPU 0
|
||||
device_map["model.embed_tokens"] = 0
|
||||
device_map["model.norm"] = 0
|
||||
device_map["lm_head"] = 0
|
||||
|
||||
print(f" Created device_map with {len(device_map)} entries")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map=device_map,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded successfully")
|
||||
|
||||
pattern = check_gpu_memory()
|
||||
print(f"\n Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_4():
|
||||
"""Test 4: Load to CPU, then move to GPU manually"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 4: Load to CPU, then move to GPU")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Loading bf16 model to CPU...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map="cpu",
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to CPU")
|
||||
|
||||
# Count params on CPU
|
||||
cpu_params = sum(p.numel() for p in model.parameters() if p.device.type == 'cpu')
|
||||
print(f" CPU parameters: {cpu_params / 1e9:.2f}B")
|
||||
|
||||
# Move to GPU 0
|
||||
print("\n Moving to GPU 0...")
|
||||
model = model.to("cuda:0")
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern after move to GPU 0: {pattern}")
|
||||
|
||||
# Move to GPU 1
|
||||
print("\n Moving to GPU 1...")
|
||||
model = model.to("cuda:1")
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern after move to GPU 1: {pattern}")
|
||||
|
||||
return True, "LOADED_TO_CPU_THEN_GPU"
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_5():
|
||||
"""Test 5: Sequential layer loading (manual distribution)"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 5: Sequential layer loading (manual distribution)")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Loading model layer by layer...")
|
||||
|
||||
# This is a simplified version - in reality would need more complex logic
|
||||
# For now, just test if we can load to one GPU
|
||||
print(" Loading bf16 to GPU 0 only...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map={"": 0},
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to GPU 0")
|
||||
|
||||
pattern = check_gpu_memory()
|
||||
print(f"\n Pattern: {pattern}")
|
||||
|
||||
# Now try GPU 1
|
||||
torch.cuda.empty_cache()
|
||||
print("\n Loading bf16 to GPU 1 only...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map={"": 1},
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to GPU 1")
|
||||
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern: {pattern}")
|
||||
|
||||
return True, "SEQUENTIAL_LOAD"
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
return False, str(e)
|
||||
# def test_strategy_5():
|
||||
# """Test 5: Sequential layer loading (manual distribution)"""
|
||||
# print("\n" + "=" * 80)
|
||||
# print("TEST 5: Sequential layer loading (manual distribution)")
|
||||
# print("=" * 80)
|
||||
#
|
||||
# try:
|
||||
# torch.cuda.empty_cache()
|
||||
# print(" Loading model layer by layer...")
|
||||
#
|
||||
# # This is a simplified version - in reality would need more complex logic
|
||||
# # For now, just test if we can load to one GPU
|
||||
# print(" Loading bf16 to GPU 0 only...")
|
||||
# model = AutoModelForCausalLM.from_pretrained(
|
||||
# "/data/models/Ornith-1.0-35B",
|
||||
# device_map={"": 0},
|
||||
# torch_dtype=torch.bfloat16,
|
||||
# trust_remote_code=True,
|
||||
# low_cpu_mem_usage=True,
|
||||
# )
|
||||
# print(" ✓ Model loaded to GPU 0")
|
||||
#
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f"\n Pattern: {pattern}")
|
||||
#
|
||||
# # Now try GPU 1
|
||||
# torch.cuda.empty_cache()
|
||||
# print("\n Loading bf16 to GPU 1 only...")
|
||||
# model = AutoModelForCausalLM.from_pretrained(
|
||||
# "/data/models/Ornith-1.0-35B",
|
||||
# device_map={"": 1},
|
||||
# torch_dtype=torch.bfloat16,
|
||||
# trust_remote_code=True,
|
||||
# low_cpu_mem_usage=True,
|
||||
# )
|
||||
# print(" ✓ Model loaded to GPU 1")
|
||||
#
|
||||
# pattern = check_gpu_memory()
|
||||
# print(f" Pattern: {pattern}")
|
||||
#
|
||||
# return True, "SEQUENTIAL_LOAD"
|
||||
# except Exception as e:
|
||||
# print(f"\n ✗ FAILED: {e}")
|
||||
# return False, str(e)
|
||||
|
||||
def test_strategy_6():
|
||||
"""Test 6: Load bf16 to CPU, quantize with BnB, then move to GPU"""
|
||||
@@ -273,6 +273,196 @@ def test_strategy_6():
|
||||
traceback.print_exc()
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_7():
|
||||
"""Test 7: bf16 to CPU → BnB 4-bit → Use accelerate to distribute"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 7: bf16 to CPU → BnB 4-bit → accelerate device_map")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Step 1: Load bf16 model to CPU...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map="cpu",
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to CPU (~70GB)")
|
||||
|
||||
print("\n Step 2: Apply BnB 4-bit quantization...")
|
||||
from peft import prepare_model_for_kbit_training
|
||||
model = prepare_model_for_kbit_training(
|
||||
model,
|
||||
use_gradient_checkpointing=False,
|
||||
)
|
||||
print(" ✓ Model prepared for k-bit training")
|
||||
|
||||
print(" Step 3: Quantize weights to 4-bit...")
|
||||
model.quantize_4bit()
|
||||
print(" ✓ Model quantized to 4-bit (~17.5GB)")
|
||||
|
||||
print("\n Step 4: Use accelerate to distribute across GPUs...")
|
||||
from accelerate import infer_auto_device_map, init_empty_weights
|
||||
from accelerate.utils import get_balanced_memory
|
||||
|
||||
# Create device map
|
||||
device_map = infer_auto_device_map(
|
||||
model,
|
||||
max_memory={0: "15GB", 1: "15GB"},
|
||||
no_split_module_classes=["Qwen3_5MoeDecoderLayer"],
|
||||
)
|
||||
print(f" Created device_map with {len(device_map)} entries")
|
||||
|
||||
# Load model with device_map
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map=device_map,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded with device_map")
|
||||
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_8():
|
||||
"""Test 8: bf16 to CPU → BnB 4-bit → Load to GPU 0 only"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 8: bf16 to CPU → BnB 4-bit → GPU 0 only")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Step 1: Load bf16 model to CPU...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map="cpu",
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to CPU (~70GB)")
|
||||
|
||||
print("\n Step 2: Apply BnB 4-bit quantization...")
|
||||
from peft import prepare_model_for_kbit_training
|
||||
model = prepare_model_for_kbit_training(
|
||||
model,
|
||||
use_gradient_checkpointing=False,
|
||||
)
|
||||
print(" ✓ Model prepared for k-bit training")
|
||||
|
||||
print(" Step 3: Quantize weights to 4-bit...")
|
||||
model.quantize_4bit()
|
||||
print(" ✓ Model quantized to 4-bit (~17.5GB)")
|
||||
|
||||
print("\n Step 4: Move to GPU 0 only...")
|
||||
model = model.to("cuda:0")
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_9():
|
||||
"""Test 9: bf16 to CPU → BnB 4-bit (int8) → GPU"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 9: bf16 to CPU → BnB 8-bit → GPU")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Step 1: Load bf16 model to CPU...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map="cpu",
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to CPU (~70GB)")
|
||||
|
||||
print("\n Step 2: Apply BnB 8-bit quantization...")
|
||||
from peft import prepare_model_for_kbit_training
|
||||
model = prepare_model_for_kbit_training(
|
||||
model,
|
||||
use_gradient_checkpointing=False,
|
||||
)
|
||||
print(" ✓ Model prepared for k-bit training")
|
||||
|
||||
print(" Step 3: Quantize weights to 8-bit...")
|
||||
# Use int8 quantization instead of 4-bit
|
||||
from bitsandbytes.nn.modules import Params8bit
|
||||
# Note: This is a simplified version - actual int8 quantization may need different approach
|
||||
print(" ⚠ int8 quantization may not be fully implemented")
|
||||
|
||||
print("\n Step 4: Move to GPU...")
|
||||
model = model.to("cuda:0")
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False, str(e)
|
||||
|
||||
def test_strategy_10():
|
||||
"""Test 10: bf16 to CPU → FSDP → GPU (single process)"""
|
||||
print("\n" + "=" * 80)
|
||||
print("TEST 10: bf16 to CPU → FSDP (single process)")
|
||||
print("=" * 80)
|
||||
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
print(" Step 1: Load bf16 model to CPU...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"/data/models/Ornith-1.0-35B",
|
||||
device_map="cpu",
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
print(" ✓ Model loaded to CPU (~70GB)")
|
||||
|
||||
print("\n Step 2: Apply FSDP wrapping...")
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
from torch.distributed.fsdp import MixedPrecision
|
||||
|
||||
# Wrap model with FSDP
|
||||
model = FSDP(
|
||||
model,
|
||||
mixed_precision=MixedPrecision(
|
||||
param_dtype=torch.bfloat16,
|
||||
reduce_dtype=torch.float32,
|
||||
buffer_dtype=torch.float32,
|
||||
),
|
||||
auto_wrap_policy=None,
|
||||
)
|
||||
print(" ✓ Model wrapped with FSDP")
|
||||
|
||||
print("\n Step 3: Move to GPU...")
|
||||
model = model.to("cuda:0")
|
||||
pattern = check_gpu_memory()
|
||||
print(f" Pattern: {pattern}")
|
||||
return True, pattern
|
||||
except Exception as e:
|
||||
print(f"\n ✗ FAILED: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False, str(e)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 80)
|
||||
print("Testing multiple model loading strategies")
|
||||
@@ -290,12 +480,16 @@ if __name__ == "__main__":
|
||||
results = []
|
||||
|
||||
tests = [
|
||||
("Test 1: device_map=auto (no BnB)", test_strategy_1),
|
||||
("Test 2: device_map=auto + BnB 4-bit", test_strategy_2),
|
||||
("Test 3: Explicit device_map", test_strategy_3),
|
||||
("Test 4: Load to CPU then GPU", test_strategy_4),
|
||||
("Test 5: Sequential layer loading", test_strategy_5),
|
||||
# ("Test 1: device_map=auto (no BnB)", test_strategy_1),
|
||||
# ("Test 2: device_map=auto + BnB 4-bit", test_strategy_2),
|
||||
# ("Test 3: Explicit device_map", test_strategy_3),
|
||||
# ("Test 4: Load to CPU then GPU", test_strategy_4),
|
||||
# ("Test 5: Sequential layer loading", test_strategy_5),
|
||||
("Test 6: bf16 to CPU → BnB 4-bit → GPU", test_strategy_6),
|
||||
("Test 7: bf16 to CPU → BnB 4-bit → accelerate", test_strategy_7),
|
||||
("Test 8: bf16 to CPU → BnB 4-bit → GPU 0 only", test_strategy_8),
|
||||
("Test 9: bf16 to CPU → BnB 8-bit → GPU", test_strategy_9),
|
||||
("Test 10: bf16 to CPU → FSDP → GPU", test_strategy_10),
|
||||
]
|
||||
|
||||
for name, test_func in tests:
|
||||
|
||||
Reference in New Issue
Block a user