Distributed Layerwise Offload¶

This document describes distributed layerwise offload (DLO) for diffusion models. DLO keeps only a small number of DiT blocks on the accelerator and streams the remaining blocks from host memory. The distributed backend can either shard those host-side weights across an existing parallel group or keep complete rank-local block sources and avoid an additional collective.

For user-facing commands, see the distributed layerwise offloading guide and the Cosmos3 recipe.

Feature compatibility¶

Host-storage optimization and runtime compatibility are separate decisions. When direct checkpoint mmap is unavailable, DLO can still use tensors produced by the ordinary loader. "Compatibility path" below means that fallback is implemented but has less end-to-end coverage than the primary path.

Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.

See Parallelism compatibility and Request and loading constraints for the detailed contracts and validation boundaries.

Status¶

DLO is implemented for multi-device diffusion execution. The default AllGather path is the primary path for DP and SP deployments. The --dlo-no-use-allgather path streams complete blocks independently and adds no DLO weight collective.

Host storage is selected separately from the transfer protocol. The loader can produce a direct-checkpoint mmap plan for a proven-compatible runtime layout; otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on the same node can share immutable checkpoint pages when direct mmap is selected, while the ordinary-loader fallback still keeps a private runtime copy per process.

The Phase A shared-mmap support boundary is TP1. TP greater than one is an ordinary-loader compatibility path: DLO can consume the resulting TP-local tensors, but those configurations do not use checkpoint mmap and must not be used to claim shared-mmap host-memory savings.

The compatibility matrix below describes the current implementation. The unit-level guards are covered, but not every parallelism combination has a full model-and-hardware end-to-end test.

Design¶

DLO consumes the existing parallel topology¶

DLO does not create a new DP, TP, or SP topology. It reads the configured DiffusionParallelConfig and attaches offload hooks to the DiT blocks after the standard distributed groups have been initialized.

The DLO weight-sharding group is selected as follows:

1. Use the existing DP group when data_parallel_size > 1.
2. When DP is one and SP is greater than one, use the SP group.
3. Otherwise, run rank-locally without a DLO process group.

TP is deliberately not used as DLO's AllGather group. HSDP has its own parameter-sharding lifecycle and is not allowed to be sharded a second time by DLO's AllGather path.

The loader owns host-weight planning¶

Before it decides whether ordinary weight materialization can be skipped, the diffusion loader builds one HostWeightPlan. A direct-checkpoint mmap plan is accepted only when preflight proves all of the following:

• every required DiT parameter and persistent buffer has exactly one source;
• runtime names, checkpoint keys, shapes, and dtypes match;
• the runtime topology is TP1 without HSDP or online quantization; and
• every custom loader operation is represented by a loader-owned checkpoint adapter.

The exact plan object is handed to DLO. The backend does not rescan checkpoint files, repeat the capability decision, or reconstruct names from its block topology. If preflight fails, the loader materializes weights normally and DLO consumes those runtime tensors.

The plan owns only dedicated DiT component sources. If a pipeline also exposes ordinary sources for a text encoder or another non-DiT component, the loader still consumes those sources and includes their loaded names in its strict coverage check. Only the source prefixes covered by the plan skip ordinary materialization. A source that mixes DiT and non-DiT weights fails closed to the complete ordinary-loader path because it cannot be skipped safely as a unit.

This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline flags such as _supports_mmap_loading or parameter attributes for mmap-only transforms. Model-specific direct-layout knowledge, when required, lives in a checkpoint adapter beside the ordinary loader.

AllGather path¶

With the default dlo_use_allgather=True, each rank stores approximately 1 / group_size of each streamable block in pinned host memory. The next block's shard is copied to a device buffer and reconstructed with all_gather_into_tensor on a communication stream while the current block is executing.

Compute:    [Block N]             [Block N+1]          [Block N+2]
H2D:                      [shard N+1]           [shard N+2]
AllGather:                [full N+1]             [full N+2]
Buffers:    [current slot]       [prefetch slot]       [current slot]

The backend uses two shared device buffers, so accelerator weight residency is bounded by the largest streamed blocks rather than the complete model.

When direct checkpoint mmap is selected, the checkpoint mappings are only the source used to prepare each rank's persistent shard. They can be closed after shard preparation. Across the AllGather group, those private shards total approximately one runtime model copy.

An effective DLO group size of one performs no collective, even when dlo_use_allgather=True; it follows the rank-local transfer path described below.

When DP is greater than one, the engine can process one request per DP rank in the same denoising wave. Because AllGather is a collective, all participating requests must take the same execution path at every denoising step.

Rank-local path without DLO AllGather¶

With --dlo-no-use-allgather, DLO forces its internal offload shard size to one and streams complete blocks using H2D copies only. The host backing may be either a loader-approved checkpoint mapping or ordinary runtime tensors.

For direct mmap, each process retains immutable safetensors views and uses two bounded pinned host staging slots. Processes on the same node that map the same files share physical checkpoint pages through the OS page cache. This removes the persistent private full-model copy per pure-DP process, but each process still packs and transfers every complete block. Sharing is node-local; each node has its own page cache.

When direct mmap preflight fails, the regular model loader remains responsible for preparing each rank's weights, including TP-local tensors or HSDP-managed parameters. In that fallback, each pure-DP process keeps a private full runtime copy.

This mode means:

• DP still provides independent replicas, but DLO does not shard weights across DP ranks.
• SP still performs its normal activation/attention collectives, but DLO does not shard weights across SP ranks.
• TP/HSDP/SP collectives, if configured, are not disabled by this flag; only DLO's additional weight AllGather is disabled.
• Pure DP deployments share one checkpoint-backed copy per node when direct mmap is selected; the ordinary-loader fallback keeps one private runtime copy per rank.
• The scheduler does not require a synchronized DP request wave for DLO.

Component allocator-cache retention¶

MiniMax-H3 DLO keeps one bounded PyTorch allocator cache across its staged encoder, DiT, and VAE component boundaries so allocations released by one component can be reused by the next. The DiT prefetch path and its two shared device buffers are unchanged.

After a component is offloaded, cached-but-unallocated memory is retained only while it is at most 25% of device capacity and at least 5% of device capacity remains physically free. Crossing either bound releases the allocator cache. Missing allocator telemetry also releases it conservatively. Component or staging failure forces a release, and an out-of-memory allocation gets one retry after release. This is a component-local policy rather than a global empty_cache override, so the executor's unconditional shutdown cleanup is preserved.

Retaining allocator blocks trades higher reserved and peak device memory for lower lifecycle latency. The benefit is device- and workload-dependent, and a smaller or externally loaded device may cross a bound and return to the ordinary per-component cache-release behavior. Current performance validation covers MiniMax-H3 no-AllGather DLO on CUDA. Other MiniMax-H3 DLO transfer topologies and platforms receive the same bounded policy but are not claimed performance targets; other DLO models remain unchanged unless they explicitly adopt it.

Final-layout Host Weight Runtime consumer¶

Status: PR2 landed in PR #6486. PR3 adds registered direct H2D plus generic HWR publication/source-digest optimizations without changing AllGather behavior.

Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies only to no-AllGather DLO. Enable it with --host-weight-runtime-mode preferred (or required) and --host-weight-runtime-root <node-local-root>.

The modes express operator fallback policy, not different artifact formats:

• preferred consumes an exact hit, but on a miss it allows canonical loading followed by post-load publication. It is the normal population path.
• required consumes the same exact artifact but fails startup on a miss or unusable artifact. It never invokes canonical DiT fallback or post-load publication and therefore cannot populate an empty store.

PR2 has no separate prewarm command. Operators populate one matching producer cohort per node-local storage domain in preferred mode, then restart the same model revision and parallel layout with required. TP coordinates have distinct identities, while equivalent DP replicas share them.

The current producer/consumer boundary is the model-declared final-layout BF16 contract (MiniMax H3 and black-forest-labs/FLUX.2-klein-4B). The FLUX.2-klein contract covers both transformer block stacks, constructor-stable packed QKV mapping state, BF16 parameters, and any persistent loader-owned beta/eps buffers. The shared HWR machinery supports ordinary-loader final layouts for TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence in this change is TP1-only; its TP2/SP layouts rely on the existing per-coordinate identity mechanics and remain unmeasured. Online quantization, HSDP, LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode disabled, DLO-disabled, and DLO AllGather configurations stop before source identity or store construction and retain the existing checkpoint-mmap or ordinary-loader path.

Eligibility is decided before constructing HWR, resolving canonical sources, hashing identity inputs, probing a filesystem, or emitting an HWR observer event:

DLO disabled / HWR disabled / DLO AllGather
  -> zero final-layout HWR interaction
  -> preserve current checkpoint-mmap or ordinary-loader behavior

When final-layout HWR is explicitly enabled for no-AllGather DLO, runtime outcomes are authoritative:

• LOCAL_HIT: plan and commit an exact restoration, then transfer the lease transactionally to DLO.
• CANONICAL_FALLBACK: bypass checkpoint mmap, canonically materialize and finalize the DiT, then call publish_after_load() for a future startup. The current startup retains its canonical tensors.
• FAILED: fail startup. Preferred mode does not reinterpret nonretryable identity, configuration, or compatibility failures as misses.
• Required mode cannot bootstrap an empty store through a POST_LOAD_ONLY producer.

Checkpoint mmap remains an unchanged control path whenever final-layout HWR is not selected.

Local source identity and publication avoid repeated full-payload passes while preserving the same semantic contract:

• Canonical files without a trusted immutable Hub blob identity are hashed in parallel. Their content IDs are cached per storage domain under the HWR root and reused only when the path, inode, size, timestamps, symlink target, and cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or cache-I/O failure falls back to hashing the canonical file directly.
• Producers that emit payload bytes in canonical storage-key order may declare ordered=True. The filesystem writer then computes file and tensor SHA256 values during the write and overlaps payload fsync with later producer work. Producers without that guarantee retain unordered writes and use a parallel readback checksum fallback.
• Manifest and READY publication still wait for every payload fsync and checksum to complete. These optimizations do not weaken artifact validation or change lease ownership.

Registered mmap transport¶

After DLO takes an HWR lease, the no-AllGather backend may register the lease's complete immutable mapped ranges under the existing pin_cpu_memory policy. Registration is transport state: it does not change artifact identity, the store, tensor ownership, or H2D payload. On success, each tensor view copies directly into the existing rotating HBM block buffers and the two private host staging slots are not allocated.

--dlo-host-registration-limit-gib is an optional per-worker preflight ceiling over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned memory policy, unsupported platform/capability, over-budget mapping, or fully rolled-back registration error selects the existing two-slot staging path. A partial registration that cannot be rolled back aborts startup because closing the lease would unmap memory still owned by the platform.

Direct checkpoint mmap remains unchanged and continues to use staging. It may require loader-owned per-block transforms, while the HWR artifact already contains final runtime bytes. DLO AllGather never receives an HWR final-layout lease and therefore never enters this registration path.

Pre-service transaction boundary¶

The startup transaction does not end at restore commit. A lease-backed model is disposable until backend setup and initial prefetch complete:

UNRESOLVED
  -> LEASE_OWNED_BY_LOADER
  -> RESTORE_PLANNED
  -> COMMIT_STARTED
  -> CARRIER_OWNED
  -> BACKEND_OWNED
  -> BACKEND_READY
  -> IN_SERVICE

Failure before COMMIT_STARTED closes the loader-owned lease. Preferred mode may canonically load using the untouched skeleton; required mode fails.

Failure from COMMIT_STARTED through BACKEND_READY must:

1. synchronize or otherwise quiesce partial backend and initial-prefetch work;
2. release hook, staging, restore-plan, and model references;
3. close the lease through its current owner;
4. discard the restored model; and
5. in preferred mode, construct a fresh canonical model while bypassing HWR lookup, HWR publication, and checkpoint mmap for that recovery attempt.

Required mode fails instead of constructing the fresh fallback. Once the model enters service, failures follow normal runtime handling rather than startup fallback.

Finalization phases¶

Cold canonical loading separates byte-changing work from shared runtime-state finalization:

Cold-only byte-changing:
  casting, reordering, packing, quantization, generated scales,
  model-specific weight transforms, and calibration that mutates
  parameters or persistent buffers

Shared cold/warm non-byte:
  validation, hook installation, eval state, bookkeeping,
  and non-persistent runtime state

The warm path must snapshot restored tensor bytes and backing pointers before shared finalization and prove both are unchanged afterward.

Lease ownership¶

Lease transfer is single-owner and single-take:

loader owns
  -> restore plan borrows
  -> commit succeeds
  -> runner carrier takes once
  -> backend takes once before asynchronous work
  -> backend drains pending H2D work
  -> backend releases hooks, staging, and model references
  -> backend unregisters every mapped range
  -> backend closes lease

The runner carrier is process-local, rejects serialization, and rejects a duplicate take(). If the carrier still owns the lease, runner cleanup closes it. Once the backend takes ownership, only backend abort or teardown may drain work and close it.

The implementation keeps the final-layout lease through backend setup and initial prefetch. Successful registration copies directly from immutable HWR views; otherwise the backend uses the existing two bounded rank-local staging slots. Transactional cleanup drains device work, removes source references, unregisters mappings, and only then closes the lease. An unregistration failure retains both registration and lease for retry/process teardown. Preferred mode then uses the runner's fresh canonical retry; required mode propagates the failure.

PR2 and PR3 promotion gates¶

• Warm hit performs zero ordinary DiT materialization and zero producer calls.
• Shared warm finalization changes neither restored bytes nor backing pointers.
• Preferred FAILED outcomes do not silently fall back.
• Planning failure leaves the skeleton untouched.
• Any pre-service failure after commit begins discards the restored model; preferred mode uses a fresh canonical model and required mode fails.
• Mixed components load normally and no required tensor remains on meta.
• Disabled mode and AllGather emit zero final-layout HWR interaction.
• The lease carrier rejects duplicate take and serialization.
• Backend setup or prefetch failure drains asynchronous work before lease close.
• Registration is all-or-nothing; rollback/unregistration failure never closes a mapping still owned by the platform.
• Registered HWR transport bypasses host staging while preserving output and H2D payload; unsupported registration retains the bounded staging path.
• Compatible checkpoint mmap remains an unchanged benchmark and control path.
• Prewarm uses one matching producer cohort per TP/SP topology and storage domain.
• Benchmarks compare ordinary loading, cold publication, and warm staging with output parity, startup latency, aggregate PSS, page-cache state, HBM, and H2D payload.

Parallelism compatibility¶

Combined dimensions¶

• DP + SP: DLO uses the DP group for weight sharding when DP is greater than one; SP continues to use its own sequence-parallel group. If DP is one, the SP group becomes DLO's sharding group in AllGather mode.
• DP + TP/SP without AllGather: standard model loading defines the rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight collective. When the model declares the final-layout contract, HWR keys the reusable artifact by TP rank/size and SP layout.
• HSDP + SP: the general parallel configuration permits HSDP over SP, but DLO must use --dlo-no-use-allgather. HSDP remains responsible for weight materialization and synchronization.
• HSDP + DP or TP: rejected independently by the diffusion parallel configuration.

Request and loading constraints¶

AllGather DP multi-concurrency requires:

• explicit num_inference_steps;
• the same num_inference_steps for all requests in a wave; and
• identical request arguments that affect the collective execution path.

The no-AllGather path does not impose these DLO-specific synchronized-wave requirements.

Direct checkpoint mmap can back either transfer path. It is currently limited to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the ordinary loader. Per-tensor online FP8, online INT8, and online MXFP8 linears can use DLO AllGather after the ordinary loader finalizes their runtime weights and scales; DLO then shards and reconstructs those tensors with their recorded layouts (they all keep plain 1-byte dtypes over ordinary strided views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view finalized tensor, all covered by the same physical-order packing as online FP8). Other online methods must use --dlo-no-use-allgather or disable online quantization until their runtime layouts are validated.

The Host Weight Runtime representation, publication, and no-AllGather consumer contracts are merged; see RFC #6414, PR #6445, and PR #6486. Registered direct H2D is an optional transport layer over that merged lease contract.

Validation coverage¶

Current source-level validation includes:

• HSDP + DLO + AllGather rejection;
• HSDP + DLO without AllGather acceptance at configuration level;
• loader preflight fallback for TP, HSDP, online quantization, unknown custom loaders, missing keys, and shape/dtype mismatches;
• ordinary-loader fallback for per-tensor online FP8, online INT8, and online MXFP8 linears followed by DLO sharding of finalized weights and scales;
• exact loader-to-backend plan transfer and ordinary-loader fallback;
• ordered publication hashing, overlapped durability, unordered parallel checksum fallback, and node-local source-digest reuse/invalidation;
• rank-local mmap source retention, bounded two-slot staging, and adapter transforms without parameter-side flags;
• read-only registration preflight, direct source-to-device copies, safe fallback, partial rollback, retryable unregistration, and lease ordering;
• resident-layer requests requiring no-AllGather;
• DP request-wave validation for denoising-step compatibility;
• bounded component allocator-cache retention, conservative/forced release, OOM retry, and immutable encoder non-block staging;
• sharding, double-buffer, AllGather-size, and heterogeneous-block regression tests.

B300 parallel-topology smoke matrix¶

A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed, CUDNN attention backend, 256x256 output, two denoising steps, and dlo_resident_layers=0. The TP2 rows used DiT DP2xTP2 with the text encoder and VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap or shared-mmap host-memory savings.

Within each topology, the AllGather and no-AllGather video and audio outputs were byte-identical. All four runs completed without an ERROR or traceback and released their device allocations. For DP4xTP1, no-AllGather direct mmap reduced total PSS by 24.22 GiB (11.4%) and Private_Dirty from 211.33 to 125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above AllGather, because DP replicas did not share checkpoint-backed runtime weights. This is a functional and memory smoke test, not a production-quality performance or output-quality benchmark.

Host-memory measurement¶

A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the ordinary-loader fallback with direct mmap. Both runs used DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a 256x256 four-second request. The ordinary-loader workers were sampled after initialization. The mmap workers were sampled after one completed request, so the checkpoint working set had been faulted into the page cache; this is the more conservative point for mmap.

The values below come from /proc/<worker>/smaps_rollup and include the whole worker, not only the DiT. The stable rank-to-rank difference comes from other pipeline components, so each worker should be compared with the same worker in the other storage mode.

The direct-mmap workers each reported 62.45 GiB Shared_Clean but only 31.20 GiB Pss_File, which is the proportional charge expected when the same resident checkpoint pages are mapped by two workers. Private_Dirty also fell from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker 1, a reduction of about 97–98 GiB per worker. RSS understates this benefit because it counts a shared physical page in every process that maps it; summed PSS is the appropriate node-memory comparison.

The highest-value missing coverage is broader end-to-end numerical and lifecycle comparison against ordinary layerwise offload for DP+SP, HSDP+SP+no-AllGather, and TP greater than one across additional models and target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not change the Phase A direct-mmap TP1 support boundary.

Recommendations¶

• Use DP + DLO AllGather for the supported throughput and host-memory scaling path.
• Use SP + DLO AllGather for long-sequence workloads when DP concurrency is not the goal.
• Use no-AllGather when independent replica execution is required. TP1 direct-mmap deployments can share checkpoint pages per node; other layouts retain the ordinary loader's private host memory behavior and are outside the Phase A shared-mmap support scope.
• Prefer HSDP alone for production HSDP deployments until the combined HSDP + DLO no-AllGather path has broader end-to-end coverage.