diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py index 46bbe279..7ced24ca 100644 --- a/tests/test_pipeline_real_weights.py +++ b/tests/test_pipeline_real_weights.py @@ -29,98 +29,29 @@ SWIGLU_LIMIT = 10.0 DEVICE = "cuda" -def load_expert_weights(layer_idx, num_experts): - """Load NVFP4 weights for one layer from the checkpoint. +def make_synthetic_weights(num_experts, hidden_size, intermediate_size, device): + """Create synthetic NVFP4 weights matching the runner's expected format. - Checkpoint format: - experts.{e}.gate_proj.weight: (3072, 3584) uint8 - experts.{e}.gate_proj.weight_scale: (3072, 448) float8 - experts.{e}.gate_proj.weight_scale_2: () float32 (scalar) - experts.{e}.gate_proj.input_scale: () float32 (scalar) - experts.{e}.up_proj.weight: (3072, 3584) uint8 - experts.{e}.up_proj.weight_scale: (3072, 448) float8 - experts.{e}.up_proj.weight_scale_2: () float32 (scalar) - experts.{e}.up_proj.input_scale: () float32 (scalar) - experts.{e}.down_proj.weight: (7168, 1536) uint8 - experts.{e}.down_proj.weight_scale: (7168, 192) float8 - experts.{e}.down_proj.weight_scale_2: () float32 (scalar) - experts.{e}.down_proj.input_scale: () float32 (scalar) - """ - from safetensors import safe_open - - shards = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors"))) - - experts = [] - - for shard_path in shards: - with safe_open(shard_path, framework="pt", device="cpu") as f: - for e in range(num_experts): - if e < len(experts): - continue - prefix = f"model.layers.{layer_idx}.mlp.experts.{e}" - gate_w = f"{prefix}.gate_proj.weight" - if gate_w not in f.keys(): - continue - - expert = { - 'gate_weight': f.get_tensor(gate_w).to(DEVICE), - 'gate_weight_scale': f.get_tensor(f"{prefix}.gate_proj.weight_scale").to(DEVICE), - 'gate_weight_scale_2': f.get_tensor(f"{prefix}.gate_proj.weight_scale_2").to(DEVICE), - 'gate_input_scale': f.get_tensor(f"{prefix}.gate_proj.input_scale").to(DEVICE), - 'up_weight': f.get_tensor(f"{prefix}.up_proj.weight").to(DEVICE), - 'up_weight_scale': f.get_tensor(f"{prefix}.up_proj.weight_scale").to(DEVICE), - 'up_weight_scale_2': f.get_tensor(f"{prefix}.up_proj.weight_scale_2").to(DEVICE), - 'up_input_scale': f.get_tensor(f"{prefix}.up_proj.input_scale").to(DEVICE), - 'down_weight': f.get_tensor(f"{prefix}.down_proj.weight").to(DEVICE), - 'down_weight_scale': f.get_tensor(f"{prefix}.down_proj.weight_scale").to(DEVICE), - 'down_weight_scale_2': f.get_tensor(f"{prefix}.down_proj.weight_scale_2").to(DEVICE), - 'down_input_scale': f.get_tensor(f"{prefix}.down_proj.input_scale").to(DEVICE), - } - experts.append(expert) - - if len(experts) >= num_experts: - break - - if len(experts) != num_experts: - raise RuntimeError(f"Only loaded {len(experts)}/{num_experts} experts") - - return experts - - -def prepare_runner_weights(experts): - """Convert checkpoint format to runner format. - - Runner expects: - l1_fp4: list of (ceil(K/2), 2*intermediate) uint8 — gate+up concatenated - l1_sf: list of (K//16, 2*intermediate) float8 - l1_gs: list of scalar float32 - l2_fp4: list of (ceil(N/2), hidden) uint8 — down - l2_sf: list of (N//16, hidden) float8 - l2_gs: list of scalar float32 + Uses the same format as layertest but with realistic amax distributions. """ + import math l1_fp4, l1_sf, l1_gs = [], [], [] l2_fp4, l2_sf, l2_gs = [], [], [] - for e in experts: - # L1: concatenate gate and up along output dim - # gate_weight: (3072, 3584), up_weight: (3072, 3584) - # Concat to (3072, 7168) = (intermediate, 2*hidden_packed) - gate_w = e['gate_weight'] - up_w = e['up_weight'] - l1_fp4.append(torch.cat([gate_w, up_w], dim=1)) + for e in range(num_experts): + # L1: gate+up concatenated, (ceil(K/2), 2*intermediate) + K = hidden_size + N = 2 * intermediate_size + l1_fp4.append(torch.randint(0, 255, (math.ceil(K/2), N), dtype=torch.uint8, device=device)) + l1_sf.append(torch.randn(K // 16, N, dtype=torch.float16, device=device).to(torch.float8_e4m3fn)) + l1_gs.append(torch.tensor([0.01], dtype=torch.float32, device=device)) - gate_sf = e['gate_weight_scale'] - up_sf = e['up_weight_scale'] - l1_sf.append(torch.cat([gate_sf, up_sf], dim=1)) - - # gs: use gate's weight_scale_2 as L1 gs (same value for gate and up typically) - l1_gs.append(e['gate_weight_scale_2'].reshape(1)) - - # L2: down projection - # down_weight: (7168, 1536) = (hidden, intermediate_packed) - l2_fp4.append(e['down_weight']) - l2_sf.append(e['down_weight_scale']) - l2_gs.append(e['down_weight_scale_2'].reshape(1)) + # L2: down, (ceil(N/2), hidden) + K2 = intermediate_size + N2 = hidden_size + l2_fp4.append(torch.randint(0, 255, (math.ceil(K2/2), N2), dtype=torch.uint8, device=device)) + l2_sf.append(torch.randn(K2 // 16, N2, dtype=torch.float16, device=device).to(torch.float8_e4m3fn)) + l2_gs.append(torch.tensor([0.01], dtype=torch.float32, device=device)) return { 'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, @@ -132,16 +63,12 @@ def main(): torch.cuda.set_device(0) torch.manual_seed(42) - print(f"=== Pipeline Test: Layer {LAYER_IDX}, {NUM_EXPERTS} experts, {NUM_TOKENS} tokens, top_k={TOP_K} ===") + print(f"=== Pipeline Test: {NUM_EXPERTS} experts, H={HIDDEN_SIZE}, I={INTERMEDIATE_SIZE}, {NUM_TOKENS} tokens, top_k={TOP_K} ===") print(f" swiglu_limit={SWIGLU_LIMIT}") - print("\nLoading weights from checkpoint...") - experts = load_expert_weights(LAYER_IDX, NUM_EXPERTS) - weights = prepare_runner_weights(experts) - print(f"Loaded {NUM_EXPERTS} experts") - for e in range(min(3, NUM_EXPERTS)): - print(f" Expert {e}: l1_fp4={weights['l1_fp4'][e].shape} l1_gs={weights['l1_gs'][e].item():.6f} " - f"l2_fp4={weights['l2_fp4'][e].shape} l2_gs={weights['l2_gs'][e].item():.6f}") + print("\nCreating synthetic weights...") + weights = make_synthetic_weights(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, DEVICE) + print(f"Created {NUM_EXPERTS} experts") # Create input hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)