diff --git a/zig/src/core/fixtures/glm53_flash_nvfp4/config.json b/zig/src/core/fixtures/glm53_flash_nvfp4/config.json new file mode 100644 index 0000000000..d0e8e75de4 --- /dev/null +++ b/zig/src/core/fixtures/glm53_flash_nvfp4/config.json @@ -0,0 +1,469 @@ +{ + "architectures": [ + "Glm5NextForConditionalGeneration" + ], + "dtype": "bfloat16", + "image_end_token_id": 154831, + "image_start_token_id": 154830, + "image_token_id": 154854, + "model_type": "glm5_next", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "base_model_ep_plan": { + "layers.*.mlp.experts": "torch_a2a_experts", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.gate": "torch_a2a_router" + }, + "bos_token_id": null, + "dtype": "bfloat16", + "eos_token_id": [ + 154820, + 154827, + 154829 + ], + "first_k_dense_replace": 3, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 0, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_kpool": 4, + "index_kpool_always_select_tail": true, + "index_kpool_compress": true, + "index_n_heads": 32, + "index_share_for_mtp_iteration": true, + "index_topk": 2048, + "indexer_rope_interleave": true, + "indexer_types": [ + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full" + ], + "initializer_range": 0.02, + "intermediate_size": 12288, + "kv_lora_rank": 512, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention" + ], + "linear_attn_config": { + "full_attn_layers": [ + 3, + 7, + 11, + 15, + 19, + 23, + 27, + 31, + 35, + 39, + 43 + ], + "gate_lower_bound": -5.0, + "head_dim": 128, + "kda_layers": [ + 0, + 1, + 2, + 4, + 5, + 6, + 8, + 9, + 10, + 12, + 13, + 14, + 16, + 17, + 18, + 20, + 21, + 22, + 24, + 25, + 26, + 28, + 29, + 30, + 32, + 33, + 34, + 36, + 37, + 38, + 40, + 41, + 42, + 44 + ], + "num_heads": 64, + "short_conv_kernel_size": 4 + }, + "linear_conv_kernel_dim": 4, + "linear_head_dim": 128, + "linear_lower_bound": -5.0, + "linear_num_heads": 64, + "max_position_embeddings": 1048576, + "mhc": true, + "mla_use_nope": true, + "mlp_layer_types": [ + "dense", + "dense", + "dense", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse" + ], + "model_type": "glm5_next_text", + "moe_intermediate_size": 2048, + "moe_router_dtype": "float32", + "n_group": 1, + "n_routed_experts": 288, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 8, + "num_hidden_layers": 45, + "num_key_value_heads": 64, + "num_nextn_predict_layers": 1, + "output_router_logits": false, + "pad_token_id": 154820, + "q_lora_rank": 1536, + "qk_head_dim": 256, + "qk_nope_head_dim": 256, + "qk_rope_head_dim": 0, + "rms_norm_eps": 1e-05, + "routed_scaling_factor": 2.5, + "router_aux_loss_coef": 0.001, + "scoring_func": "sigmoid", + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "topk_group": 1, + "topk_method": "noaux_tc", + "use_cache": true, + "v_head_dim": 256, + "vocab_size": 154880 + }, + "tie_word_embeddings": false, + "transformers_version": "5.16.1", + "video_end_token_id": 154833, + "video_start_token_id": 154832, + "video_token_id": 154855, + "vision_config": { + "attention_bias": true, + "attention_dropout": 0.0, + "depth": 24, + "dtype": "bfloat16", + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 448, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "glm5_next_vision", + "num_heads": 16, + "out_hidden_size": 4096, + "patch_size": 14, + "projection_intermediate_size": 10240, + "rms_norm_eps": 1e-05, + "spatial_merge_size": 2, + "swiglu_limit": 10.0, + "temporal_patch_size": 2 + }, + "quantization_config": { + "config_groups": { + "group_0": { + "input_activations": { + "dynamic": false, + "num_bits": 4, + "type": "float", + "group_size": 16 + }, + "weights": { + "dynamic": false, + "num_bits": 4, + "type": "float", + "group_size": 16 + }, + "targets": [ + "Linear" + ] + } + }, + "ignore": [ + "lm_head", + "model.language_model.embed_tokens", + "model.language_model.layers.0.self_attn*", + "model.language_model.layers.1.self_attn*", + "model.language_model.layers.10.mlp.gate", + "model.language_model.layers.10.mlp.shared_experts*", + "model.language_model.layers.10.self_attn*", + "model.language_model.layers.11.mlp.gate", + "model.language_model.layers.11.mlp.shared_experts*", + "model.language_model.layers.11.self_attn*", + "model.language_model.layers.12.mlp.gate", + "model.language_model.layers.12.mlp.shared_experts*", + "model.language_model.layers.12.self_attn*", + "model.language_model.layers.13.mlp.gate", + "model.language_model.layers.13.mlp.shared_experts*", + "model.language_model.layers.13.self_attn*", + "model.language_model.layers.14.mlp.gate", + "model.language_model.layers.14.mlp.shared_experts*", + "model.language_model.layers.14.self_attn*", + "model.language_model.layers.15.mlp.gate", + "model.language_model.layers.15.mlp.shared_experts*", + "model.language_model.layers.15.self_attn*", + "model.language_model.layers.16.mlp.gate", + "model.language_model.layers.16.mlp.shared_experts*", + "model.language_model.layers.16.self_attn*", + "model.language_model.layers.17.mlp.gate", + "model.language_model.layers.17.mlp.shared_experts*", + "model.language_model.layers.17.self_attn*", + "model.language_model.layers.18.mlp.gate", + "model.language_model.layers.18.mlp.shared_experts*", + "model.language_model.layers.18.self_attn*", + "model.language_model.layers.19.mlp.gate", + "model.language_model.layers.19.mlp.shared_experts*", + "model.language_model.layers.19.self_attn*", + "model.language_model.layers.2.self_attn*", + "model.language_model.layers.20.mlp.gate", + "model.language_model.layers.20.mlp.shared_experts*", + "model.language_model.layers.20.self_attn*", + "model.language_model.layers.21.mlp.gate", + "model.language_model.layers.21.mlp.shared_experts*", + "model.language_model.layers.21.self_attn*", + "model.language_model.layers.22.mlp.gate", + "model.language_model.layers.22.mlp.shared_experts*", + "model.language_model.layers.22.self_attn*", + "model.language_model.layers.23.mlp.gate", + "model.language_model.layers.23.mlp.shared_experts*", + "model.language_model.layers.23.self_attn*", + "model.language_model.layers.24.mlp.gate", + "model.language_model.layers.24.mlp.shared_experts*", + "model.language_model.layers.24.self_attn*", + "model.language_model.layers.25.mlp.gate", + "model.language_model.layers.25.mlp.shared_experts*", + "model.language_model.layers.25.self_attn*", + "model.language_model.layers.26.mlp.gate", + "model.language_model.layers.26.mlp.shared_experts*", + "model.language_model.layers.26.self_attn*", + "model.language_model.layers.27.mlp.gate", + "model.language_model.layers.27.mlp.shared_experts*", + "model.language_model.layers.27.self_attn*", + "model.language_model.layers.28.mlp.gate", + "model.language_model.layers.28.mlp.shared_experts*", + "model.language_model.layers.28.self_attn*", + "model.language_model.layers.29.mlp.gate", + "model.language_model.layers.29.mlp.shared_experts*", + "model.language_model.layers.29.self_attn*", + "model.language_model.layers.3.mlp.gate", + "model.language_model.layers.3.mlp.shared_experts*", + "model.language_model.layers.3.self_attn*", + "model.language_model.layers.30.mlp.gate", + "model.language_model.layers.30.mlp.shared_experts*", + "model.language_model.layers.30.self_attn*", + "model.language_model.layers.31.mlp.gate", + "model.language_model.layers.31.mlp.shared_experts*", + "model.language_model.layers.31.self_attn*", + "model.language_model.layers.32.mlp.gate", + "model.language_model.layers.32.mlp.shared_experts*", + "model.language_model.layers.32.self_attn*", + "model.language_model.layers.33.mlp.gate", + "model.language_model.layers.33.mlp.shared_experts*", + "model.language_model.layers.33.self_attn*", + "model.language_model.layers.34.mlp.gate", + "model.language_model.layers.34.mlp.shared_experts*", + "model.language_model.layers.34.self_attn*", + "model.language_model.layers.35.mlp.gate", + "model.language_model.layers.35.mlp.shared_experts*", + "model.language_model.layers.35.self_attn*", + "model.language_model.layers.36.mlp.gate", + "model.language_model.layers.36.mlp.shared_experts*", + "model.language_model.layers.36.self_attn*", + "model.language_model.layers.37.mlp.gate", + "model.language_model.layers.37.mlp.shared_experts*", + "model.language_model.layers.37.self_attn*", + "model.language_model.layers.38.mlp.gate", + "model.language_model.layers.38.mlp.shared_experts*", + "model.language_model.layers.38.self_attn*", + "model.language_model.layers.39.mlp.gate", + "model.language_model.layers.39.mlp.shared_experts*", + "model.language_model.layers.39.self_attn*", + "model.language_model.layers.4.mlp.gate", + "model.language_model.layers.4.mlp.shared_experts*", + "model.language_model.layers.4.self_attn*", + "model.language_model.layers.40.mlp.gate", + "model.language_model.layers.40.mlp.shared_experts*", + "model.language_model.layers.40.self_attn*", + "model.language_model.layers.41.mlp.gate", + "model.language_model.layers.41.mlp.shared_experts*", + "model.language_model.layers.41.self_attn*", + "model.language_model.layers.42.mlp.gate", + "model.language_model.layers.42.mlp.shared_experts*", + "model.language_model.layers.42.self_attn*", + "model.language_model.layers.43.mlp.gate", + "model.language_model.layers.43.mlp.shared_experts*", + "model.language_model.layers.43.self_attn*", + "model.language_model.layers.44.mlp.gate", + "model.language_model.layers.44.mlp.shared_experts*", + "model.language_model.layers.44.self_attn*", + "model.language_model.layers.5.mlp.gate", + "model.language_model.layers.5.mlp.shared_experts*", + "model.language_model.layers.5.self_attn*", + "model.language_model.layers.6.mlp.gate", + "model.language_model.layers.6.mlp.shared_experts*", + "model.language_model.layers.6.self_attn*", + "model.language_model.layers.7.mlp.gate", + "model.language_model.layers.7.mlp.shared_experts*", + "model.language_model.layers.7.self_attn*", + "model.language_model.layers.8.mlp.gate", + "model.language_model.layers.8.mlp.shared_experts*", + "model.language_model.layers.8.self_attn*", + "model.language_model.layers.9.mlp.gate", + "model.language_model.layers.9.mlp.shared_experts*", + "model.language_model.layers.9.self_attn*", + "model.visual*" + ], + "quant_algo": "NVFP4", + "kv_cache_scheme": { + "dynamic": false, + "num_bits": 8, + "type": "float" + }, + "producer": { + "name": "modelopt", + "version": "0.47.0.dev393+ga4bc45b30.d20260828" + }, + "quant_method": "modelopt" + } +} \ No newline at end of file diff --git a/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz b/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz new file mode 100644 index 0000000000..db569985ec Binary files /dev/null and b/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz differ diff --git a/zig/src/core/modelopt.zig b/zig/src/core/modelopt.zig new file mode 100644 index 0000000000..da37ad7b5e --- /dev/null +++ b/zig/src/core/modelopt.zig @@ -0,0 +1,110 @@ +//! ModelOpt NVFP4 config check (groups of 16); tensors say which linears are NVFP4, not ignore (423acf37's BF16 MTP). + +const std = @import("std"); + +/// The NVFP4 group size ModelOpt writes for this format. +pub const nvfp4_group = 16; + +/// Why the check refused the checkpoint, kept for the caller's log line (tests read it instead). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + /// Keep the reason, cut short if it is long. + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +/// config.json's quantization_config (top level, else text_config's) checked as ModelOpt NVFP4. +pub fn admit(gpa: std.mem.Allocator, config_json: []const u8, why: *Why) !void { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const root = try std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), config_json, .{}); + const q = block(root) orelse { + why.set("config.json has no ModelOpt quantization_config", .{}); + return error.NotModelOpt; + }; + const algo = text(q.get("quant_algo")) orelse ""; + if (!std.ascii.eqlIgnoreCase(algo, "NVFP4")) { + why.set("ModelOpt {s}: this engine reads ModelOpt NVFP4 (nvidia/GLM-5.3-Flash-NVFP4's layout)", .{if (algo.len > 0) algo else "without quant_algo"}); + return error.UnsupportedQuant; + } + const groups = object(q.get("config_groups")) orelse { + why.set("ModelOpt NVFP4 without config_groups", .{}); + return error.UnsupportedQuant; + }; + var it = groups.iterator(); + while (it.next()) |g| { + const w = if (object(g.value_ptr.*)) |o| object(o.get("weights")) else null; + const bits = if (w) |o| int(o.get("num_bits")) else null; + const size = if (w) |o| int(o.get("group_size")) else null; + const kind = if (w) |o| text(o.get("type")) orelse "" else ""; + if (bits != 4 or size != nvfp4_group or !std.mem.eql(u8, kind, "float")) { + why.set("ModelOpt NVFP4 group {s}: weights {?d}-bit {s} in groups of {?d}, expected 4-bit float in groups of 16", .{ g.key_ptr.*, bits, kind, size }); + return error.UnsupportedQuant; + } + } +} + +fn block(root: std.json.Value) ?std.json.ObjectMap { + const top = object(root) orelse return null; + for ([_]?std.json.ObjectMap{ top, object(top.get("text_config")) }) |source| { + const q = object((source orelse continue).get("quantization_config")) orelse continue; + if (std.ascii.eqlIgnoreCase(text(q.get("quant_method")) orelse "", "modelopt")) return q; + } + return null; +} + +fn object(v: ?std.json.Value) ?std.json.ObjectMap { + const x = v orelse return null; + return if (x == .object) x.object else null; +} + +fn text(v: ?std.json.Value) ?[]const u8 { + const x = v orelse return null; + return if (x == .string) x.string else null; +} + +fn int(v: ?std.json.Value) ?i64 { + const x = v orelse return null; + return if (x == .integer) x.integer else null; +} + +test "nvidia/GLM-5.3-Flash-NVFP4's config is admitted" { + var why: Why = .{}; + try admit(std.testing.allocator, @embedFile("fixtures/glm53_flash_nvfp4/config.json"), &why); +} + +/// admit() must refuse `json` with `err`, its reason naming `word`. +fn expectRefused(json: []const u8, err: anyerror, word: []const u8) !void { + var why: Why = .{}; + try std.testing.expectError(err, admit(std.testing.allocator, json, &why)); + if (std.mem.indexOf(u8, why.text(), word) == null) { + std.debug.print("refusal \"{s}\" does not name \"{s}\"\n", .{ why.text(), word }); + return error.TestUnexpectedResult; + } +} + +test "other ModelOpt formats, and configs without ModelOpt, are refused by name" { + try expectRefused("{\"text_config\": {}}", error.NotModelOpt, "no ModelOpt"); + try expectRefused("{\"quantization_config\": {\"quant_method\": \"exl3\"}}", error.NotModelOpt, "no ModelOpt"); + try expectRefused("{\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"FP8\"}}", error.UnsupportedQuant, "ModelOpt FP8"); + // MIXED_PRECISION (the AXL re-pack) comes with 3.x + try expectRefused("{\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"MIXED_PRECISION\"}}", error.UnsupportedQuant, "MIXED_PRECISION"); + try expectRefused( + \\{"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "config_groups": + \\ {"g": {"weights": {"num_bits": 4, "type": "float", "group_size": 32}}}}} + , error.UnsupportedQuant, "groups of 32"); + var why: Why = .{}; + try admit(std.testing.allocator, + \\{"text_config": {"quantization_config": {"quant_method": "modelopt", "quant_algo": "nvfp4", "config_groups": + \\ {"g": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}}}}}} + , &why); +} diff --git a/zig/src/core/nvfp4.zig b/zig/src/core/nvfp4.zig new file mode 100644 index 0000000000..ed6948a003 --- /dev/null +++ b/zig/src/core/nvfp4.zig @@ -0,0 +1,123 @@ +//! NVFP4 as ModelOpt stores it: e2m1 codes two to a byte, an e4m3 scale per 16 inputs, an fp32 scale per tensor. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); +const DType = ckpt.DType; + +/// Inputs sharing one e4m3 scale. +pub const group = quant.nvfp4.group; + +/// A stored tensor's dtype and shape, as a header lists it. +pub const Meta = struct { + dtype: DType, + shape: []const usize, +}; + +pub const Dims = struct { n: usize, k: usize }; + +/// The linear's [n, k] from codes U8 [n, k/2], scales F8_E4M3 [n, k/16] and fp32 scalars; else error.BadNvfp4. +pub fn dims(weight: Meta, scale: Meta, scale2: Meta, input_scale: ?Meta) !Dims { + if (weight.dtype != .u8 or weight.shape.len != 2) return error.BadNvfp4; + if (scale.dtype != .f8_e4m3 or scale.shape.len != 2) return error.BadNvfp4; + const n = weight.shape[0]; + const k = 2 * weight.shape[1]; + if (k % group != 0 or scale.shape[0] != n or scale.shape[1] * group != k) return error.BadNvfp4; + if (!scalar(scale2)) return error.BadNvfp4; + if (input_scale) |s| if (!scalar(s)) return error.BadNvfp4; + return .{ .n = n, .k = k }; +} + +fn scalar(m: Meta) bool { + return m.dtype == .f32 and (m.shape.len == 0 or (m.shape.len == 1 and m.shape[0] == 1)); +} + +/// Whether `module` (a linear's name without ".weight") is stored as NVFP4: its tensors say, not config.json. +pub fn stored(ck: *const ckpt.Checkpoint, module: []const u8) bool { + var buf: [256]u8 = undefined; + const name = std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}) catch return false; + for (ck.files.items) |*f| if (f.names.contains(name)) return true; + return false; +} + +/// One NVFP4 linear's stored bytes, as the checkpoint maps them. +pub const Linear = struct { + n: usize, + k: usize, + codes: []const u8, // [n, k/2]: input 2j in byte j's low nibble + scales: []const u8, // [n, k/16] e4m3 + global: f32, // weight_scale_2 + + /// `module`'s four tensors, marked as used; input_scale is the W4A4 math's static scale, unread by W4A16. + pub fn read(ck: *ckpt.Checkpoint, module: []const u8) !Linear { + var buf: [256]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const s = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const a = try ck.get(try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module})); + const d = dims(meta(&w), meta(&s), meta(&g), meta(&a)) catch |e| { + std.log.err("{s}: not ModelOpt NVFP4 ({t} {any}, {t} {any})", .{ module, w.dtype, w.shape[0..w.rank], s.dtype, s.shape[0..s.rank] }); + return e; + }; + return .{ .n = d.n, .k = d.k, .codes = w.bytes, .scales = s.bytes, .global = std.mem.bytesToValue(f32, g.bytes[0..4]) }; + } +}; + +fn meta(t: *const ckpt.Tensor) Meta { + return .{ .dtype = t.dtype, .shape = t.shape[0..t.rank] }; +} + +test "dims refuse what is not ModelOpt NVFP4" { + const w: Meta = .{ .dtype = .u8, .shape = &.{ 4, 16 } }; + const s: Meta = .{ .dtype = .f8_e4m3, .shape = &.{ 4, 2 } }; + const g: Meta = .{ .dtype = .f32, .shape = &.{} }; + const d = try dims(w, s, g, g); + try std.testing.expectEqual(Dims{ .n = 4, .k = 32 }, d); + try std.testing.expectError(error.BadNvfp4, dims(.{ .dtype = .bf16, .shape = &.{ 4, 16 } }, s, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, .{ .dtype = .f8_e4m3, .shape = &.{ 4, 3 } }, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, .{ .dtype = .u8, .shape = &.{ 4, 2 } }, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, s, .{ .dtype = .f32, .shape = &.{2} }, null)); +} + +test "a linear read from a checkpoint folder decodes as core/quant's reference" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + // n = 2, k = 32: codes 2 x 16 bytes, scales 2 x 2, two fp32 scalars + const header = + \\{"l.weight":{"dtype":"U8","shape":[2,16],"data_offsets":[0,32]}, + \\ "l.weight_scale":{"dtype":"F8_E4M3","shape":[2,2],"data_offsets":[32,36]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[36,40]}, + \\ "l.input_scale":{"dtype":"F32","shape":[],"data_offsets":[40,44]}} + ; + var data: [8 + header.len + 44]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + const body = data[8 + header.len ..]; + for (body[0..32], 0..) |*b, i| b.* = @intCast((i * 7 + 3) & 0xff); // every nibble value appears + @memcpy(body[32..36], &[_]u8{ 0x38, 0x40, 0x30, 0xb8 }); // 1.0, 2.0, 0.5, -1.0 + @memcpy(body[36..40], std.mem.asBytes(&@as(f32, 0.25))); + @memcpy(body[40..44], std.mem.asBytes(&@as(f32, 3.0))); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &data }); + const dir = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer gpa.free(dir); + var ck = try ckpt.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + try std.testing.expect(stored(&ck, "l")); + try std.testing.expect(!stored(&ck, "m")); + const l = try Linear.read(&ck, "l"); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + try std.testing.expectEqual(@as(f32, 0.25), l.global); + var out: [64]f64 = undefined; + try quant.nvfp4.reference(.{ + .weight = .{ .dtype = .u8, .rank = 2, .shape = .{ l.n, l.k / 2, 1, 1, 1 }, .bytes = l.codes }, + .scale = .{ .dtype = .f8_e4m3, .rank = 2, .shape = .{ l.n, l.k / group, 1, 1, 1 }, .bytes = l.scales }, + .global = &[_]f32{l.global}, + }, &out); + // byte 0 is 3: input 0 is code 3 (1.5), input 1 code 0; inputs 16.. take the second scale (2.0) + try std.testing.expectEqual(@as(f64, 1.5 * 1.0 * 0.25), out[0]); + try std.testing.expectEqual(@as(f64, 0.0), out[1]); + try std.testing.expectEqual(quant.nvfp4.e2m1(body[8] & 0xf) * 2.0 * 0.25, out[16]); + try std.testing.expectEqual(quant.nvfp4.e2m1(body[16 + 15] >> 4) * -1.0 * 0.25, out[32 + 31]); // row 1, scale -1.0 +} diff --git a/zig/src/core/nvfp4_fixture_test.zig b/zig/src/core/nvfp4_fixture_test.zig new file mode 100644 index 0000000000..9f4bfa3694 --- /dev/null +++ b/zig/src/core/nvfp4_fixture_test.zig @@ -0,0 +1,69 @@ +//! nvidia/GLM-5.3-Flash-NVFP4's whole header: every linear is NVFP4 or BF16 as nvfp4.dims reads it. + +const std = @import("std"); +const nvfp4 = @import("nvfp4.zig"); +const DType = @import("safetensors.zig").DType; + +const Meta = nvfp4.Meta; + +/// name -> [dtype, shape] from the gzip'd fixture, kept in `arena`. +fn headers(arena: std.mem.Allocator) !std.json.ObjectMap { + var in: std.Io.Reader = .fixed(@embedFile("fixtures/glm53_flash_nvfp4/headers.json.gz")); + var gz: std.compress.flate.Decompress = .init(&in, .gzip, &.{}); + var out: std.Io.Writer.Allocating = .init(arena); + _ = try gz.reader.streamRemaining(&out.writer); + const v = try std.json.parseFromSliceLeaky(std.json.Value, arena, out.written(), .{}); + return v.object; +} + +fn meta(arena: std.mem.Allocator, h: std.json.ObjectMap, name: []const u8) !?Meta { + const v = h.get(name) orelse return null; + const pair = v.array.items; + const dims = pair[1].array.items; + const shape = try arena.alloc(usize, dims.len); + for (dims, shape) |d, *s| s.* = @intCast(d.integer); + return .{ .dtype = DType.parse(pair[0].string) orelse return error.UnsupportedDType, .shape = shape }; +} + +fn at(arena: std.mem.Allocator, h: std.json.ObjectMap, module: []const u8, suffix: []const u8) !?Meta { + return meta(arena, h, try std.fmt.allocPrint(arena, "{s}.{s}", .{ module, suffix })); +} + +test "every NVFP4 linear of the checkpoint reads as ModelOpt stores it, and the MTP layer stays BF16" { + var state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer state.deinit(); + const a = state.allocator(); + const h = try headers(a); + try std.testing.expectEqual(@as(usize, 147661), h.count()); + var quantized: usize = 0; + var experts: usize = 0; + var it = h.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + if (std.mem.endsWith(u8, name, ".weight_scale")) { + // no scale without its codes + try std.testing.expect((try at(a, h, name[0 .. name.len - ".weight_scale".len], "weight")) != null); + continue; + } + if (!std.mem.endsWith(u8, name, ".weight")) continue; + const module = name[0 .. name.len - ".weight".len]; + const scale = try at(a, h, module, "weight_scale") orelse continue; + const d = try nvfp4.dims((try meta(a, h, name)).?, scale, (try at(a, h, module, "weight_scale_2")).?, try at(a, h, module, "input_scale")); + quantized += 1; + try std.testing.expect(std.mem.indexOf(u8, module, ".mlp.") != null); + try std.testing.expect(std.mem.indexOf(u8, module, ".layers.45.") == null); + const down = std.mem.endsWith(u8, module, "down_proj"); + if (std.mem.indexOf(u8, module, ".experts.") != null) { + experts += 1; + try std.testing.expectEqual(if (down) nvfp4.Dims{ .n = 4096, .k = 2048 } else nvfp4.Dims{ .n = 2048, .k = 4096 }, d); + } else { + try std.testing.expectEqual(if (down) nvfp4.Dims{ .n = 4096, .k = 12288 } else nvfp4.Dims{ .n = 12288, .k = 4096 }, d); + } + } + try std.testing.expectEqual(@as(usize, 42 * 288 * 3 + 3 * 3), quantized); + try std.testing.expectEqual(@as(usize, 42 * 288 * 3), experts); + // the MTP layer's routed experts are BF16 as stored (a family may re-pack them for drafting) + const mtp = (try meta(a, h, "model.language_model.layers.45.mlp.experts.0.gate_proj.weight")).?; + try std.testing.expectEqual(DType.bf16, mtp.dtype); + try std.testing.expect((try meta(a, h, "model.language_model.layers.45.mlp.experts.0.gate_proj.weight_scale")) == null); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..b3d7cfbc48 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -18,6 +18,8 @@ pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); pub const act_ref = @import("act_ref.zig"); +pub const nvfp4 = @import("nvfp4.zig"); +pub const modelopt = @import("modelopt.zig"); test { _ = safetensors; @@ -37,4 +39,7 @@ test { _ = draft_depth; _ = @import("copy_index.zig"); _ = ids_json; + _ = nvfp4; + _ = modelopt; + _ = @import("nvfp4_fixture_test.zig"); }