refactor: polish roadmap APIs and memory export defaults

Agent-Logs-Url: https://github.com/codewithdark-git/QuantLLM/sessions/fc1f2077-187e-4757-9927-2c60ef187666

Co-authored-by: codewithdark-git <144595403+codewithdark-git@users.noreply.github.com>

Author: Copilot

docs/guide/gguf-export.md

@@ -209,6 +209,9 @@ model.export("gguf", "model.gguf", quantization="Q4_K_M")
 For very large models:
 
 ```python
+# Note: previous `streaming=True` guidance is superseded by `chunked_conversion=True`.
+# If you previously used `streaming=True`, replace it with `chunked_conversion=True` (streaming has no effect here).
+
 # Use lower quantization
 model.export("gguf", "model.Q3_K_M.gguf", quantization="Q3_K_M")
 

quantllm/core/memory.py

@@ -12,7 +12,7 @@
 """
 
 import gc
-from typing import Optional, Dict, Any, List, Union, Callable, Iterable
+from typing import Optional, Dict, Any, List, Union, Callable
 from contextlib import contextmanager
 from collections import OrderedDict
 import torch
@@ -191,16 +191,16 @@ def memory_optimized_tensor_order(
     state_dict: Dict[str, torch.Tensor],
     *,
     prioritize_large_tensors: bool = True,
-) -> "OrderedDict[str, torch.Tensor]":
+) -> OrderedDict[str, torch.Tensor]:
     """
     Return an ordered state dict to reduce peak memory pressure during serialization.
     
-    By default, larger tensors are emitted first to reduce long-lived allocator pressure
+    By default, tensors are sorted by total byte size (numel * element_size),
+    with larger tensors emitted first to reduce long-lived allocator pressure
     in shard-based writes on very large checkpoints.
     """
-    items: Iterable = state_dict.items()
     sorted_items = sorted(
-        items,
+        state_dict.items(),
         key=lambda kv: kv[1].numel() * kv[1].element_size(),
         reverse=prioritize_large_tensors,
     )

quantllm/core/turbo_model.py

@@ -23,6 +23,9 @@
 from ..utils import logger, print_header, print_success, print_error, print_info, print_warning, QuantLLMProgress
 from transformers.utils.logging import disable_progress_bar as disable_hf_progress_bar
 from datasets.utils.logging import disable_progress_bar as disable_ds_progress_bar
+from .memory import memory_optimized_tensor_order
+
+DEFAULT_CHUNKED_SHARD_SIZE = "2GB"
 
 
 class TurboModel:
@@ -1127,6 +1130,11 @@ def _export_gguf(
         output_path: str, 
         quantization: Optional[str] = None,
         fast_mode: bool = False,
+        chunked_conversion: bool = False,
+        max_shard_size: Optional[str] = None,
+        smart_tensor_ordering: bool = False,
+        disk_offloading: bool = False,
+        disk_offload_dir: Optional[str] = None,
         **kwargs
     ) -> str:
         """
@@ -1156,11 +1164,9 @@ def _export_gguf(
 
         start_time = time.time()
 
-        chunked_conversion = bool(kwargs.pop("chunked_conversion", False))
-        max_shard_size = kwargs.pop("max_shard_size", "2GB" if chunked_conversion else "50GB")
-        smart_tensor_ordering = bool(kwargs.pop("smart_tensor_ordering", False))
-        disk_offloading = bool(kwargs.pop("disk_offloading", False))
-        disk_offload_dir = kwargs.pop("disk_offload_dir", None)
+        effective_shard_size = max_shard_size or (
+            DEFAULT_CHUNKED_SHARD_SIZE if chunked_conversion else None
+        )
 
         quant_type = quantization or self.config.quant_type or "q4_k_m"
         quant_type_upper = quant_type.upper()
@@ -1175,9 +1181,10 @@ def _export_gguf(
             if fast_mode:
                 print_info("Fast mode enabled")
             if chunked_conversion:
-                print_info(f"Chunked conversion enabled (max_shard_size={max_shard_size})")
+                print_info(f"Chunked conversion enabled (max_shard_size={effective_shard_size})")
             if smart_tensor_ordering:
                 print_info("Smart tensor ordering enabled")
+                print_warning("Smart tensor ordering may temporarily materialize a full state dict in memory.")
             if disk_offloading:
                 print_info(f"Disk offloading enabled ({disk_offload_dir or 'system temp'})")
 
@@ -1220,11 +1227,11 @@ def _export_gguf(
                 task = progress.add_task("Saving model weights...", total=None)
                 save_kwargs = {
                     "safe_serialization": True,
-                    "max_shard_size": max_shard_size,
                 }
+                if effective_shard_size:
+                    save_kwargs["max_shard_size"] = effective_shard_size
 
                 if smart_tensor_ordering:
-                    from .memory import memory_optimized_tensor_order
                     save_kwargs["state_dict"] = memory_optimized_tensor_order(model_to_save.state_dict())
 
                 try:

quantllm/kernels/triton/quantized_linear.py

@@ -76,8 +76,10 @@ def int4_matmul(
         scales: Per-column scales, shape [1, out_features] or [in_features/group, out_features]
         bias: Optional bias [out_features]
     """
-    zeros = torch.zeros_like(scales)
-    group_size = qweight.shape[0] if scales.shape[0] == 1 else max(qweight.shape[0] // scales.shape[0], 1)
+    # Per-column case uses [1, N] zeros; grouped quantization uses zeros shaped like scales.
+    is_per_column = scales.shape[0] == 1
+    zeros = scales.new_zeros((1, scales.shape[1])) if is_per_column else scales.new_zeros(scales.shape)
+    group_size = qweight.shape[0] if is_per_column else max(qweight.shape[0] // scales.shape[0], 1)
     return fused_dequant_matmul(
         x=x,
         qweight=qweight,
@@ -525,7 +527,7 @@ def extra_repr(self) -> str:
         )
 
 
-TRITON_QUANT_KERNELS: Dict[str, Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]] = {
+triton_quantizers: Dict[str, Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]] = {
     "q4_0": triton_q4_0_quantize,
     "q8_0": triton_q8_0_quantize,
 }