distributed_writer#

No-gather distributed HF checkpoint writer for FSDP2.

Every tensor is sent to exactly one owner rank over batched point-to-point sends and written from there, so no rank ever materializes the whole model and no unit is replicated across ranks. Used for every expert-parallel model; a model sharded by FSDP2 alone opts in with MODELOPT_FSDP2_EXPORT=nogather.

Functions

distributed_save_hf_checkpoint

No-gather distributed HF export.
distributed_save_hf_checkpoint(model, export_dir, maxbound, kv_cache_format, max_shard_size='10GB', is_modelopt_qlora=False, extra_state_dict=None)#

No-gather distributed HF export.

Thin wrapper: hf_ptq exports under torch.inference_mode(), so run the entire write with inference mode DISABLED via a nested inference_mode(False) context – the .detach() in get_model_state_dict and the version- counted ops in the write path both reject inference tensors. The context form reliably restores normal mode for this scope (a decorator did not, under the caller’s already-active inference_mode). no_grad too: inference_mode(False) re-ENABLES autograd, so a bare clone()/detach() on an inference tensor would try to set up grad and read its (missing) version counter – no_grad avoids that while the clones still land as normal tensors because inference mode is off.

Parameters:

  • model (Module)

  • export_dir (str | Path)

  • maxbound (float)

  • kv_cache_format (str | None)

  • max_shard_size (str | int)

  • is_modelopt_qlora (bool)

  • extra_state_dict (dict[str, Tensor] | None)

Return type:

None