Update Wall-X to 1.1.0 (#104)

wall_x/trainer/utils/__init__.py

@@ -0,0 +1,6 @@
+"""Trainer utility helpers."""
+
+from .data import move_batch_to_device
+from .diagnostics import log_gpu_memory
+
+__all__ = ["move_batch_to_device", "log_gpu_memory"]

wall_x/trainer/utils/data.py

@@ -0,0 +1,46 @@
+"""Data-related utility functions used by the trainer main loop."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def move_batch_to_device(
+    batch: Any,
+    device: torch.device,
+    *,
+    non_blocking: bool = True,
+) -> Any:
+    """Move every tensor in ``batch`` to ``device``, recursing into dict/list.
+
+    Returns a new structure with the same shape; the input ``batch`` is not
+    mutated. dict / list containers are rebuilt; tensors are moved via
+    ``.to(device, non_blocking=...)``; everything else is passed through
+    by reference.
+
+    Parameters
+    ----------
+    batch : Any
+        Typically a dict produced by the dataloader, but recursion accepts
+        dict / list / tensor / arbitrary leaves.
+    device : torch.device
+        Target device, typically ``self.device`` on the trainer.
+    non_blocking : bool
+        Whether to use pinned-memory async copy. Default True because the
+        trainer uses pinned loaders; pass False if the dataloader hasn't
+        pinned memory.
+    """
+    if isinstance(batch, dict):
+        return {
+            k: move_batch_to_device(v, device, non_blocking=non_blocking)
+            for k, v in batch.items()
+        }
+    if isinstance(batch, list):
+        return [
+            move_batch_to_device(v, device, non_blocking=non_blocking) for v in batch
+        ]
+    if isinstance(batch, torch.Tensor):
+        return batch.to(device, non_blocking=non_blocking)
+    return batch

wall_x/trainer/utils/diagnostics.py

@@ -0,0 +1,40 @@
+"""Training-run diagnostic helpers (CUDA memory, etc.)."""
+
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+import torch
+
+
+def log_gpu_memory(
+    device: torch.device,
+    rank: int,
+    *,
+    stage: str = "",
+    log_fn: Optional[Callable] = None,
+) -> None:
+    """Log per-rank GPU memory (allocated / reserved / total) via ``log_fn``.
+
+    Calls ``torch.cuda.synchronize`` on ``device`` so the numbers reflect
+    the actual post-op usage, not pending work. If ``log_fn`` is None this
+    is a no-op.
+    """
+    if log_fn is None:
+        return
+    torch.cuda.synchronize()
+    allocated = torch.cuda.memory_allocated(device) / 1024**3
+    reserved = torch.cuda.memory_reserved(device) / 1024**3
+    peak_allocated = torch.cuda.max_memory_allocated(device) / 1024**3
+    peak_reserved = torch.cuda.max_memory_reserved(device) / 1024**3
+    total = torch.cuda.get_device_properties(device).total_memory / 1024**3
+    tag = f"[{stage}] " if stage else ""
+    log_fn(
+        f"{tag}GPU memory rank{rank} "
+        f"| allocated {allocated:.2f} GiB"
+        f" | reserved {reserved:.2f} GiB"
+        f" | peak_allocated {peak_allocated:.2f} GiB"
+        f" | peak_reserved {peak_reserved:.2f} GiB"
+        f" | total {total:.2f} GiB",
+        main_process_only=False,
+    )