unified_export_hf_streaming#

Streaming HF checkpoint export for disk/CPU-offloaded models.

Kept apart from unified_export_hf so the resident exporter cannot drift back into being offload-aware: the only edge between them is the dispatch in export_hf_checkpoint, which imports _export_transformers_checkpoint_streaming() lazily to keep the dependency acyclic.

Functions

collect_export_tensors

Pack each export unit and return this rank's (key, CPU tensor) pairs.
collect_export_tensors(model, dtype, is_modelopt_qlora, *, owner)#

Pack each export unit and return this rank’s (key, CPU tensor) pairs.

Every rank must call this: each unit’s gather is a collective that only completes once every rank arrives. owner decides who keeps a unit – "share" deals them round-robin so each rank keeps roughly 1/world of the model, "rank0" gives every unit to rank 0, which then holds the whole thing and must have room for it.

All the gathers finish before this returns, so the caller can write or postprocess without stalling anyone. Returning a list rather than a generator is what guarantees that.

Parameters:

  • model (Module)

  • dtype (dtype)

  • is_modelopt_qlora (bool)

  • owner (str)

Return type:

list[tuple[str, Tensor]]