[Diffusion] Register HWR mmap for direct DLO H2D (#6591)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
This commit is contained in:
Hongsheng Liuandhsliuustc0106 authored and GitHub committed 2026-08-25 23:30:11 +08:00
1 parent 5a8b81ebeb
commit c608910d20
28 files changed
+1869 -153

No files matched your search

@@ -65,6 +65,7 @@ omni = Omni(
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
## Host-weight loading
@@ -178,12 +179,48 @@ returning to `required`.
On a cold start, the canonical loader remains authoritative and publishes a
validated final-layout artifact for later workers. A warm start restores the
DiT final tensors without ordinary DiT materialization, then DLO streams them
through the same two bounded host staging slots. The artifact identity includes
the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
layouts do not alias one another. Publication failure remains separate from a
`preferred` serving startup; the next `required` startup provides the explicit
artifact-availability check.
DiT final tensors without ordinary DiT materialization. DLO then attempts to
register the immutable mapping for direct H2D; if registration is unavailable,
it streams through the same two bounded host staging slots. The artifact
identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
and distinct SP layouts do not alias one another. Publication failure remains
separate from a `preferred` serving startup; the next `required` startup
provides the explicit artifact-availability check.
For local canonical checkpoints, the first eligible worker may hash source
shards to establish immutable identity. HWR caches those digests in the same
node-local storage domain and validates file metadata before reuse, so later
workers normally avoid repeating that read. Cold BF16 publication also hashes
ordered payloads as they are written and overlaps payload durability work with
later shards; this changes startup work only, not artifact contents or runtime
H2D behavior.
#### Registered direct H2D
Registration is attempted automatically only for an eligible warm HWR hit when
`pin_cpu_memory` is enabled. A successful path is:
```text
shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
```
It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
the two host staging slots. It does not reduce the H2D payload or make GPU
kernels access host memory directly.
CUDA requires read-only host-registration capability for the immutable HWR
mapping. Unsupported capability, a positive registration limit smaller than
the complete page-aligned mapping, or a safely rolled-back registration error
falls back to bounded staging. Programmatic configurations can set
`pin_cpu_memory=False` to disable registration explicitly. A successful
registration locks the complete mapped range in host memory for that worker's
lifetime, so use
`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
required.
Shutdown drains H2D work, releases hook/source references, unregisters every
range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
retain their existing transfer paths.
When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
not resolve HWR sources or construct its store.