• Stefy Lanza (nextime / spora )'s avatar
    longcat: INT8 weights that can be built AND loaded (v0.2.66) · 999efe6c
    Stefy Lanza (nextime / spora ) authored
    Building the INT8 DiT needed ~42 GB because it materialised the bf16 model
    first; loading it needed ~55 GB for the same reason, plus every shard
    accumulated into one dict and then copied in. Both took the host down.
    
    Both now stream. The quantiser reads tensors one at a time through safe_open
    and flushes 4 GB shards, enumerating the Linear layers on a meta-device
    skeleton instead of a real model. _load_quantized_base_dit builds its skeleton
    on meta too, then assigns each shard into it as it is read.
    
    Measured on the 13.6B DiT: 51 GB bf16 -> 13.6 GB INT8, ~4 GB peak to write,
    and the load is 20 s at 0.6 GB RSS (safetensors mmap) with all 437 quantised
    Linears present and nothing stranded on meta.
    
    A checkpoint that leaves tensors on the meta device is now an error naming
    them, rather than a model that fails later somewhere unrelated.
    Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
    Claude-Session: https://claude.ai/code/session_01QXD9yKKFZHAB9t1LfWQWwz
    999efe6c
longcat_quantize.py 11.2 KB