diff --git a/.gitignore b/.gitignore index ccfb0d2..7f879b6 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ __pycache__/ *.pyc uv.lock output/ +venv/ diff --git a/configs/calib_glm4_7.toml b/configs/calib_glm4_7.toml new file mode 100644 index 0000000..adf7c08 --- /dev/null +++ b/configs/calib_glm4_7.toml @@ -0,0 +1,9 @@ +# Calibration datasets for GLM-4.7. + +[[dataset]] +path = "text/agentic_coding_calib_v3.jsonl" +max_len = 4096 + +[[dataset]] +path = "text/deep_calib.jsonl" +max_len = 4096 diff --git a/models/__init__.py b/models/__init__.py index cc9a6ca..504090a 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -1,3 +1,4 @@ +from .glm4_7 import Glm4_7Config from .glm5 import Glm5Config from .minimax_m25 import MinimaxM25Config from .qwen3_5_122b import Qwen35_122BConfig @@ -5,6 +6,7 @@ from .qwen3_5_moe_noshared import Qwen35MoeNoSharedConfig _CONFIGS = { + "glm4_7": Glm4_7Config, "glm5": Glm5Config, "minimax_m25": MinimaxM25Config, "qwen3_5_122b": Qwen35_122BConfig, diff --git a/models/glm4_7.py b/models/glm4_7.py new file mode 100644 index 0000000..9dc854d --- /dev/null +++ b/models/glm4_7.py @@ -0,0 +1,20 @@ +from .base import ModelQuantConfig + + +Glm4_7Config = ModelQuantConfig( + model_id="zai-org/GLM-4.7", + trust_remote_code=True, + streaming=True, + extra_quant_overrides={ + "*mlp.gate*": {"enable": False}, + "*mtp*": {"enable": False}, + }, +) + + +def _register_moe(): + from moe_registry import register_glm4_7_moe_for_quantization + register_glm4_7_moe_for_quantization() + + +Glm4_7Config.register_moe = _register_moe diff --git a/moe_registry.py b/moe_registry.py index 0501fe3..c19b5b5 100644 --- a/moe_registry.py +++ b/moe_registry.py @@ -2,6 +2,7 @@ Supports: - MiniMaxM2 (built-in transformers): MiniMaxM2Experts + - GLM-4 MoE (glm4_moe): Glm4MoeMoE + Glm4MoeNaiveMoe - GLM-5 / glm_moe_dsa (remote code): GlmMoeDsaMoE + GlmMoeDsaNaiveMoe - Qwen3.5 MoE: Qwen3_5MoeSparseMoeBlock + Qwen3_5MoeExperts @@ -119,6 +120,43 @@ def forward(self, hidden_states, top_k_index, top_k_weights): return final_hidden_states +# --------------------------------------------------------------------------- +# GLM-4 MoE (glm4_moe). +# --------------------------------------------------------------------------- + +class _QuantGlm4MoeMoE(_QuantSparseMoe): + @property + def num_experts(self): + return self.n_routed_experts + + def forward(self, hidden_states): + return super(_QuantSparseMoe, self).forward(hidden_states) + + +def register_glm4_7_moe_for_quantization(): + """Register GLM-4 MoE (glm4_moe) classes with modelopt.""" + try: + from transformers.models.glm4_moe.modeling_glm4_moe import ( + Glm4MoeMoE, + Glm4MoeNaiveMoe, + ) + except ImportError: + print("⚠ glm4_moe not in transformers, skipping GLM-4 registration") + return + + if QuantModuleRegistry.get(Glm4MoeMoE) is None: + QuantModuleRegistry.register( + {Glm4MoeMoE: "Glm4MoeMoE"} + )(_QuantGlm4MoeMoE) + + if QuantModuleRegistry.get(Glm4MoeNaiveMoe) is None: + QuantModuleRegistry.register( + {Glm4MoeNaiveMoe: "Glm4MoeNaiveMoe"} + )(_QuantFusedExperts) + + print("✓ Registered GLM-4 MoE for quantization") + + # --------------------------------------------------------------------------- # GLM-5 (glm_moe_dsa). # --------------------------------------------------------------------------- diff --git a/quantize.py b/quantize.py index f37923a..103b401 100644 --- a/quantize.py +++ b/quantize.py @@ -51,6 +51,8 @@ help="Load amax checkpoint and resume calibration from where it left off.") parser.add_argument("--resume-batch", type=int, default=0, help="Skip batches before this number (1-indexed). Use with --resume-amax.") +parser.add_argument("--microbatch", type=int, default=None, + help="Split calibration batches into microbatches of this size to limit GPU 0 activation memory.") parser.add_argument("--calib-method", default="max", choices=["max", "quantile"], help="Calibration algorithm. 'quantile' uses P2 streaming quantile estimation.") parser.add_argument("--save-quantiles", type=str, default=None, diff --git a/scripts/quantize_glm4_7.sh b/scripts/quantize_glm4_7.sh new file mode 100755 index 0000000..7867797 --- /dev/null +++ b/scripts/quantize_glm4_7.sh @@ -0,0 +1,19 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")/.." +source venv/bin/activate + +export SAFETENSORS_FAST_GPU=1 +export PYTORCH_ALLOC_CONF=expandable_segments:True + +python quantize.py \ + --model glm4_7 \ + --model-id /nvme/models/safetensors/GLM-4.7/ \ + --export-dir /nvme/models/temp/GLM-4.7-NVFP4 \ + --calib-config configs/calib_glm4_7.toml \ + --cpu-capacity 256GiB \ + --streaming \ + --floor-amaxes \ + --save-amax /nvme/models/temp/GLM-4.7-NVFP4/amaxes.safetensors \ + --batch-tokens 1572864 \ diff --git a/streaming_loader.py b/streaming_loader.py index 47512c3..27b5c51 100644 --- a/streaming_loader.py +++ b/streaming_loader.py @@ -87,12 +87,13 @@ def __init__( self.num_gpus = torch.cuda.device_count() if gpu_capacity_gib is None: - # Auto-detect from first GPU - self.gpu_capacity_gib = ( - torch.cuda.get_device_properties(0).total_memory / 1024**3 - ) + # Auto-detect per-GPU capacity + self.gpu_capacities_gib = [ + torch.cuda.get_device_properties(i).total_memory / 1024**3 + for i in range(self.num_gpus) + ] else: - self.gpu_capacity_gib = gpu_capacity_gib + self.gpu_capacities_gib = [gpu_capacity_gib] * self.num_gpus self.snapshot_dir = _resolve_snapshot_dir(model_id) index_path = os.path.join(self.snapshot_dir, "model.safetensors.index.json") @@ -198,10 +199,6 @@ def _compute_layer_sizes(self, model): def _compute_storage_map(self, layer_sizes): """Assign each layer to a storage device: cuda:1-N, cpu, or meta (disk).""" storage_map = {} - # Reserve headroom per storage GPU for CUDA context, driver memory, - # and PyTorch allocator fragmentation from loading many small tensors. - gpu_headroom = 4.0 * 1024**3 - gpu_capacity = self.gpu_capacity_gib * 1024**3 - gpu_headroom cpu_capacity = self.cpu_capacity_gib * 1024**3 # GPU 0 is reserved for execution — start packing from GPU 1 @@ -214,6 +211,10 @@ def _compute_storage_map(self, layer_sizes): placed = False while current_gpu < self.num_gpus: + # Reserve ~4% headroom for CUDA context, driver memory, + # and PyTorch allocator fragmentation. + gpu_total = self.gpu_capacities_gib[current_gpu] * 1024**3 + gpu_capacity = gpu_total * 0.96 if gpu_used + size <= gpu_capacity: storage_map[i] = f"cuda:{current_gpu}" gpu_used += size