nvflare.app_opt.pt.lazy_tensor_dict module

PT lazy tensor references used by tensor disk offload.

When enable_tensor_disk_offload=True, incoming streamed tensor payloads are written to temporary safetensors files instead of being fully deserialized into memory. LazyTensorDict maps item IDs to on-disk files, and _LazyRef defers loading until materialize() is called by aggregation code.

This keeps peak memory lower for large models while still allowing deterministic explicit cleanup via cleanup(), with GC as a fallback through _TempDirRef.

class LazyTensorDict(key_to_file: dict[str, tuple[str, str]], temp_dir: str)

Bases: object

Dict-like mapping of FOBS item_ids to on-disk safetensors files.

Each entry maps an item_id to a (file_path, key) pair. Tensors are loaded via safetensors safe_open (mmap) on access.

cleanup()

get(key, default=None)

items()

keys()

make_lazy_ref(key) → _LazyRef

values()