Skip to content

online_quantization_streaming_guide.md

Latest commit

 

History

History
305 lines (231 loc) · 12.1 KB

online_quantization_streaming_guide.md

File metadata and controls

305 lines (231 loc) · 12.1 KB
Raw

Streaming online quantization guide

ATOM can quantize eligible modules while their checkpoint weights are being loaded, then release their source weights early. This reduces the load-time peak GPU memory caused by source and quantized weights coexisting.

This guide covers only the streaming execution path. For target formats, layer-matching rules, and accuracy guidance, first read the online quantization guide and online quantization best practices.

1. What streaming online quantization solves

Without streaming, online quantization runs in two consecutive phases:

  1. Load all checkpoint weights onto the device.
  2. Walk the model and call process_weights_after_loading() to create the quantized weights.

Until the second phase finishes, source weights and quantized outputs may occupy device memory at the same time. For large MoE models, this temporary peak can determine whether the model starts successfully.

Streaming reduces the processing unit to one Linear or fused MoE module:

Read the module's checkpoint weights
    ↓
Confirm that all parameters for the module have arrived
    ↓
Move them to the target device and quantize
    ↓
Release source weights no longer needed by the module
    ↓
Continue loading later modules

Checkpoint reads, H2D transfers, quantization, and source-storage release can overlap. The primary goal is to reduce load-time peak memory; the serving API after startup is unchanged.

2. Activation conditions and defaults

The following environment variables control streaming:

Environment variable Default Description
ATOM_ONLINE_QUANT_STREAMING 0 Opts in to streaming online quantization.
ATOM_ONLINE_QUANT_STREAMING_HOST_STAGING 1 Assembles module weights in CPU host memory before moving them to the device.
ATOM_ONLINE_QUANT_STREAMING_THREADS 4 Number of tail workers for H2D, quantization, and source release. 0 finalizes synchronously in checkpoint-walker threads.

Streaming is disabled by default so existing online-quant loads retain their post-load behavior and full-weight quantization semantics. Set ATOM_ONLINE_QUANT_STREAMING=1 to opt in. After it is enabled, all of the following must also hold:

  1. The command provides a valid --online_quant_config.
  2. The model contains online-quant modules that support streaming.
  3. The load is not a dummy load.
  4. The source checkpoint and target format are supported by online quantization.

Setting ATOM_ONLINE_QUANT_STREAMING=1 without an --online_quant_config does not quantize the model by itself.

3. Recommended launch configuration

The following example converts an FP8-block source model to FP8 non-expert Linear layers and MXFP4 experts:

ATOM_ONLINE_QUANT_STREAMING=1 \
ATOM_ONLINE_QUANT_STREAMING_HOST_STAGING=1 \
ATOM_ONLINE_QUANT_STREAMING_THREADS=4 \
python -m atom.entrypoints.openai_server \
  --model deepseek-ai/DeepSeek-R1-0528 \
  -tp 8 \
  --online_quant_config '{
    "global_quant_config": "ptpc_fp8",
    "layer_quant_config": {"*expert*": "mxfp4"},
    "exclude_layer": ["lm_head", "*.gate.*"]
  }'

The environment variables and --online_quant_config syntax are identical when using atom.examples.simple_inference.

A/B against the non-streaming path

Unset the streaming variable, or set it explicitly to 0, to quantize after the complete model has loaded:

ATOM_ONLINE_QUANT_STREAMING=0 \
python -m atom.entrypoints.openai_server \
  ... \
  --online_quant_config '{...}'

When comparing load time or peak memory, keep the model, TP/EP configuration, quantization JSON, application caches, and checkpoint page-cache state identical. Cold storage and a warm page cache can differ by several times in load duration.

4. Host staging and buffered replay

Host staging (default when streaming is enabled, and recommended)

With ATOM_ONLINE_QUANT_STREAMING_HOST_STAGING=1:

  1. Streamable parameters initially avoid allocating their complete GPU source storage.
  2. Checkpoint arrivals write into module-level staging storage in CPU host memory.
  3. After every parameter for a module has arrived, its parameters are moved to the target device.
  4. Tail workers quantize and release source storage early. When ATOM_ONLINE_QUANT_STREAMING_THREADS=0, checkpoint-walker threads perform this finalization synchronously instead.

This mode allows the checkpoint walker to retain the concurrency configured by ATOM_LOADER_NUM_THREADS. It usually provides better loading throughput, at the cost of corresponding host-memory usage.

Buffered replay

With ATOM_ONLINE_QUANT_STREAMING_HOST_STAGING=0, loader calls are recorded and replayed on the target device when the module can be processed. The checkpoint walk is forced to one thread for ordering and state safety. Tail quantization can still use the workers configured by ATOM_ONLINE_QUANT_STREAMING_THREADS; a value of 0 finalizes inline.

This mode is primarily useful for compatibility comparisons and debugging. It is usually not the preferred performance configuration.

5. Detecting module completion

A weight loader may slice the destination, call narrow, or map expert shards before writing. The checkpoint tensor shape alone is therefore not sufficient to determine whether a module is complete.

ATOM counts the number of elements actually written through aten.copy_ inside the loading scope. Where available, fused MoE modules use a more explicit (expert, shard) region-coverage protocol, preventing duplicate arrivals from claiming a module too early.

Only after all parameters for a module are known to be complete does the streamer:

  1. Claim the module so later arrivals cannot overwrite quantized results.
  2. Submit finalization to a tail worker, or run it inline when the tail-worker count is 0.
  3. Call process_weights_after_loading().
  4. Release stale source storage.

6. Tuning the worker count

ATOM_ONLINE_QUANT_STREAMING_THREADS controls the size of the tail-worker pool.

Value Use case Cost
0 Do not create tail workers Checkpoint-walker threads quantize synchronously
1 Memory is tight, but some overlap is desirable Low parallelism
4 Recommended default Balance between throughput and in-flight memory
>4 H2D or quantization is the measured bottleneck and memory is available More in-flight storage, streams, and scheduling overhead; not guaranteed to be faster

A value of 0 does not imply strictly serial execution. With host staging enabled, multiple checkpoint-walker threads can still finalize different modules concurrently. Set ATOM_LOADER_NUM_THREADS=1 as well when validating strict ordering.

When the tail-worker count is greater than zero, a semaphore limits submitted but incomplete tail tasks to 2 × worker count, preventing an unbounded queue when loading produces work faster than quantization consumes it.

Before increasing the worker count, measure load time, GPU peak memory, host memory, and the number of fallback modules.

7. Verifying that streaming is active

Every rank logs phase timing:

[rank_0] weight load phases (including streaming online quant):
read+queue 140.71s (...) | drain 2.01s | quant drain 0.00s |
staging flush 0.00s | threads 16

Here, threads 16 is the checkpoint-walker thread count from ATOM_LOADER_NUM_THREADS, not the tail-worker count from ATOM_ONLINE_QUANT_STREAMING_THREADS. staging flush refers to the loader's expert-staging flush, not the streamer's host-staging H2D time.

The runtime also logs streaming coverage and fallback counts. The numbers below are one example, not expected values for every model:

Online-quant streaming: 444/582 eligible modules quantized during load,
138 fell back to the post-load pass (no memory saving for those).

The final loading summary includes:

Post-stream fallback and weight processing done: 0.55 seconds
Peak GPU memory during streaming weight loading and online quantization: ...
Model load done: ... (weights loaded in 145.67s)

Check the following:

  • eligible modules quantized during load is greater than zero.
  • The fallback list is consistent with the model structure.
  • There are no warnings about parameters that never loaded and were zero-filled.
  • There are no warnings about checkpoint arrivals dropped after a module was quantized.
  • Layer counts and formats in online_quant_info_*.json match the requested quantization configuration.

8. What fallback means

Not every eligible module can necessarily be quantized during checkpoint loading. The following cases stay in the standard post-load pass:

  • A child module must retain source weights until a parent combines them.
  • A custom or fused loader cannot provide reliable write coverage.
  • Arrivals are incomplete or coverage tracking becomes invalid.
  • A module needs the loader to finish before its final storage is known.

Fallback preserves compatibility, but those modules do not receive the peak-memory benefit of streaming source release. A small, stable, and explainable fallback set is usually normal. A sudden increase should prompt an inspection of recent model-loader or module post-processing changes.

9. Performance and output semantics

The primary goal of streaming online quantization is lower load-time peak memory, not lower load time. Compared with quantizing after the complete model has loaded, it adds some fixed overhead:

  1. In the default host-staging mode, checkpoint arrivals write into host staging rather than remaining directly in final device source storage. This adds CPU copies and staging allocation/zero-fill.
  2. Every loader call records actual write coverage and updates module completion under a lock.
  3. Module completion requires storage swaps, future submission, semaphore acquisition, and device-stream management when tail workers are enabled.
  4. Workers synchronize after H2D and quantization. H2D, quantization kernels, and checkpoint reads may also compete for memory bandwidth.
  5. Fallback modules pay streaming-tracking overhead and still enter the standard post-load pass.

Streaming does not reduce the required quantization computation; it attempts to overlap that work with checkpoint I/O. Consequently:

  • With cold storage or an I/O-bound load, checkpoint reads may hide quantization, making startup equal or faster.
  • With a warm page cache, lightweight quantization, or many fallbacks, little work can be hidden and management overhead appears directly as several extra seconds of load time.
  • More workers help only when H2D or quantization is the bottleneck. Too many workers increase in-flight storage, synchronization, and bandwidth contention, and can make loading slower.

Under TP, streaming quantizes the local shard on each rank directly. When matching full-weight quantization would require a TP all-gather, quantization scales and final weights can differ slightly from non-streaming or offline full-weight quantization. Run an accuracy A/B before deployment.

The final runtime format is still selected by --online_quant_config. Streaming primarily changes load scheduling and intermediate-storage lifetime, with the TP output difference noted above.

10. Troubleshooting

No streaming logs appear

Confirm that the command contains a non-empty --online_quant_config, and that the source checkpoint supports online quantization. The streaming environment variable alone does not create quantization work.

Streaming is slower than non-streaming

First confirm that both runs use the same cache state. Then:

  1. Keep host staging enabled.
  2. Benchmark worker counts of 1, 4, and higher.
  3. Check whether too many modules fall back.
  4. Determine whether the bottleneck is checkpoint I/O, H2D, or a quantization kernel.

Peak memory is still high

Reduce ATOM_ONLINE_QUANT_STREAMING_THREADS and inspect the fallback count. Many fallback modules keep their source weights until the post-load phase.

Disable streaming completely

ATOM_ONLINE_QUANT_STREAMING=0

Online quantization still runs, but only in the unified post-load pass after the complete model has loaded.