Fix view meta passthrough and skip empty logs

This commit is contained in:
ifilipis 2026-01-08 19:49:15 +02:00
parent f8caa93595
commit 1cc2a5733c
2 changed files with 26 additions and 0 deletions

View File

@ -322,6 +322,8 @@ def _log_materialization(
context: str,
):
total_bytes, cpu_bytes, gpu_bytes, meta_bytes = _summarize_module_bytes(module, refs)
if total_bytes == 0:
return
partial = meta_bytes > 0
LOGGER.info(
"%s: module=%s dest=%s load=%0.2fMB free=%0.2fMB partial=%s "

View File

@ -639,6 +639,30 @@ class _BaseViewStateDict(MutableMapping):
self._deleted.add(key)
return self.get_tensor(key)
def meta(self, key: str):
if key in self._overrides:
t = self._overrides[key]
numel = t.numel()
return SimpleNamespace(
dtype=t.dtype,
shape=tuple(t.shape),
numel=numel,
nbytes=numel * t.element_size(),
)
base_key = self._resolve_base_key(key)
if base_key is None or key in self._deleted:
raise KeyError(key)
if hasattr(self._base, "meta"):
return self._base.meta(base_key)
t = self._base[base_key]
numel = t.numel()
return SimpleNamespace(
dtype=t.dtype,
shape=tuple(t.shape),
numel=numel,
nbytes=numel * t.element_size(),
)
class DeviceViewStateDict(_BaseViewStateDict):
def __init__(