Add Memory Prints

def _mem_gb() -> float:
    """Return current CUDA memory allocated in GiB."""
    return torch.cuda.memory_allocated() / 1024**3


def _mem_detail() -> str:
    """Return allocated, cached-pool, and non-pytorch memory in GiB."""
    free, total = torch.cuda.mem_get_info()
    allocated = torch.cuda.memory_allocated()
    reserved = torch.cuda.memory_reserved()
    G = 1024**3
    cached = reserved - allocated
    non_pytorch = total - free - reserved
    return f"alloc={allocated / G:.2f} cached={cached / G:.2f} non-pt={non_pytorch / G:.2f}"

self.memory_profiling = True  # Set to True to enable per-step memory logging

# -- Per-layer activation profiling --
_layer_mem_profile = False
_layer_mem_ranks = RANKS


class _SavedTensorsProfiler:
    """Context manager that logs tensors saved by autograd, distinguishing weights from activations."""

    def __init__(self, layer_idx: int, stage_name: str, weight_data_ptrs: set):
        self._layer_idx = layer_idx
        self._stage_name = stage_name
        self._weight_data_ptrs = weight_data_ptrs
        self._log: list[str] = []
        self._act_bytes = 0
        self._wt_bytes = 0

    def _pack(self, t: torch.Tensor) -> torch.Tensor:
        nbytes = t.nelement() * t.element_size()
        is_wt = t.data_ptr() in self._weight_data_ptrs
        tag = "weight" if is_wt else "activ"
        self._log.append(
            f"    saved ({tag}): {tuple(t.shape)} {t.dtype} ({nbytes / 1024**2:.1f} MB)"
        )
        if is_wt:
            self._wt_bytes += nbytes
        else:
            self._act_bytes += nbytes
        return t

    def __enter__(self):
        self._ctx = torch.autograd.graph.saved_tensors_hooks(self._pack, lambda t: t)
        self._ctx.__enter__()
        return self

    def __exit__(self, *args):
        self._ctx.__exit__(*args)

    def print_summary(self):
        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
        if rank not in _layer_mem_ranks:
            return
        hdr = f"layer{self._layer_idx} {self._stage_name} saved tensors"
        print(f"[rank={rank}] {hdr}:", flush=True)
        for line in self._log:
            print(f"[rank={rank}] {line}", flush=True)
        print(
            f"[rank={rank}]   activ={self._act_bytes / 1024**2:.1f} MB, "
            f"weight={self._wt_bytes / 1024**2:.1f} MB, "
            f"total={len(self._log)} tensors",
            flush=True,
        )


def _lmem(label: str) -> None:
    """Print memory at a layer-internal checkpoint."""
    if not _layer_mem_profile:
        return
    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
    if rank not in _layer_mem_ranks:
        return
    torch.cuda.synchronize()
    alloc = torch.cuda.memory_allocated() / 1024**3
    print(f"[rank={rank}]   {label}: alloc={alloc:.2f}", flush=True)

def _setup_mem(label: str) -> None:
    """Print CUDA memory at a setup checkpoint."""
    if torch.distributed.get_rank() in RANKS:
        torch.cuda.synchronize()
        G = 1024**3
        alloc = torch.cuda.memory_allocated()
        reserved = torch.cuda.memory_reserved()
        free, total = torch.cuda.mem_get_info()
        cached = reserved - alloc
        non_pytorch = total - free - reserved
        print(
            f"[rank={torch.distributed.get_rank()}] setup_model | {label}: "
            f"alloc={alloc / G:.2f} cached={cached / G:.2f} non-pt={non_pytorch / G:.2f}",
            flush=True,
        )

# Probe: isolate CUDA context cost from NCCL communicator init
torch.cuda.set_device(device_id)
torch.cuda.synchronize()
_free0, _total = torch.cuda.mem_get_info()
_non_pt0 = _total - _free0
G = 1024**3

# ... existing init_process_group code ...

torch.cuda.synchronize()
_free1, _ = torch.cuda.mem_get_info()
_non_pt1 = _total - _free1
if device_id in RANKS:
    print(
        f"[rank={torch.distributed.get_rank()}] init_process_group | "
        f"cuda_ctx={_non_pt0 / G:.2f} "
        f"after_nccl_world={_non_pt1 / G:.2f} "
        f"nccl_world_cost={(_non_pt1 - _non_pt0) / G:.2f}",
        flush=True,
    )

torch.cuda.synchronize()
_free_before, _total = torch.cuda.mem_get_info()
_non_pt_before = _total - _free_before

# ... existing init_device_mesh code ...

torch.cuda.synchronize()
_free_after, _ = torch.cuda.mem_get_info()
_non_pt_after = _total - _free_after
G = 1024**3
if device_id in RANKS:
    print(
        f"[rank={distributed.rank}] init_device_mesh | "
        f"non-pt={_non_pt_after / G:.2f} "
        f"mesh_cost={(_non_pt_after - _non_pt_before) / G:.2f}",
        flush=True,
    )

_profiling = self.memory_profiling and distributed.rank in RANKS
if _profiling:
    torch.cuda.synchronize()
    _m0 = _mem_gb()
    print(
        f"[rank={distributed.rank} pp={pp_rank}] Before pipeline: {_m0:.2f} GiB | {_mem_detail()}",
        flush=True,
    )

for i in range(step_1):
    self._forward_chunk(0)

for i in range(step_1):
    if _profiling and i == 1:
        import pithtrain.pipeline.execution as _mod
        _mod._layer_mem_profile = True
    self._forward_chunk(0)
    if _profiling and i == 1:
        _mod._layer_mem_profile = False
    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step1 F0 i={i}: {_mem_gb():.2f} GiB (+{_mem_gb() - _m0:.2f}) | {_mem_detail()}",
            flush=True,
        )

if _profiling:
    torch.cuda.synchronize()
    _m1 = _mem_gb()
    print(
        f"[rank={distributed.rank} pp={pp_rank}] After Step1 ({step_1} F0): {_m1:.2f} GiB (+{_m1 - _m0:.2f}) | {_mem_detail()}",
        flush=True,
    )

for i in range(step_2):
    self._forward_chunk(0, recv=False, send=False)
    self._recv_forward(0)
    self._forward_chunk(1, send=(not self.is_last_pp_rank) or (i < step_2 - 1))
    self._send_forward(0)

for i in range(step_2):
    self._forward_chunk(0, recv=False, send=False)
    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step2 i={i} forward_chunk(0): {_mem_gb():.2f} GiB (+{_mem_gb() - _m0:.2f}) | {_mem_detail()}",
            flush=True,
        )
    self._recv_forward(0)
    if _profiling and i == 0:
        import pithtrain.pipeline.execution as _mod
        _mod._layer_mem_profile = True
    self._forward_chunk(1, send=(not self.is_last_pp_rank) or (i < step_2 - 1))
    if _profiling and i == 0:
        _mod._layer_mem_profile = False
    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step2 i={i} forward_chunk(1): {_mem_gb():.2f} GiB (+{_mem_gb() - _m0:.2f}) | {_mem_detail()}",
            flush=True,
        )
    self._send_forward(0)

if _profiling:
    torch.cuda.synchronize()
    _m2 = _mem_gb()
    print(
        f"[rank={distributed.rank} pp={pp_rank}] After Step2 ({step_2} F0F1): {_m2:.2f} GiB (+{_m2 - _m0:.2f}) | {_mem_detail()}",
        flush=True,
    )

for i in range(step_3):
    self._backward_chunk(1, enable_zb=True)
    self._recv_forward(1)
    self._weight_chunk()
    self._forward_chunk(1, recv=False)

for i in range(step_3):
    if _profiling:
        torch.cuda.synchronize()
        _ms3 = _mem_gb()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step3 i={i} before B1: {_ms3:.2f} GiB | {_mem_detail()}",
            flush=True,
        )
    self._backward_chunk(1, enable_zb=True)
    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step3 i={i} after B1: {_mem_gb():.2f} GiB (delta={_mem_gb() - _ms3:+.2f}) | {_mem_detail()}",
            flush=True,
        )
    self._recv_forward(1)
    self._weight_chunk()
    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step3 i={i} after W1: {_mem_gb():.2f} GiB (delta={_mem_gb() - _ms3:+.2f}) | {_mem_detail()}",
            flush=True,
        )
    self._forward_chunk(1, recv=False)
    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step3 i={i} after F1: {_mem_gb():.2f} GiB (delta={_mem_gb() - _ms3:+.2f}) | {_mem_detail()}",
            flush=True,
        )

if _profiling:
    torch.cuda.synchronize()
    _m3 = _mem_gb()
    print(
        f"[rank={distributed.rank} pp={pp_rank}] After Step3 ({step_3} B1W1F1): {_m3:.2f} GiB (+{_m3 - _m0:.2f}) | {_mem_detail()}",
        flush=True,
    )

    if _profiling:
        torch.cuda.synchronize()
        print(
            f"[rank={distributed.rank} pp={pp_rank}] Step4 i={i}: {_mem_gb():.2f} GiB | {_mem_detail()}",
            flush=True,
        )

if _profiling:
    torch.cuda.synchronize()
    _m4 = _mem_gb()
    print(
        f"[rank={distributed.rank} pp={pp_rank}] After Step4 ({step_4} F0B1F1B0): {_m4:.2f} GiB | {_mem_detail()}",
        flush=True,
    )

if _profiling:
    torch.cuda.synchronize()
    _m8 = _mem_gb()
    print(
        f"[rank={distributed.rank} pp={pp_rank}] After Step8 (end of pipeline): {_m8:.2f} GiB | {_mem_detail()}",
        flush=True,
    )

_do_lmem = _layer_mem_profile and layer.idx in DETAIL_LAYERS
_do_saved = _layer_mem_profile and layer.idx in DETAIL_LAYERS

_wt_ptrs: set = set()
if _do_saved:
    _wt_ptrs = {p.data_ptr() for p in layer.parameters()}

if _do_saved:
    _prof = _SavedTensorsProfiler(layer.idx, "stage1", _wt_ptrs)
    with _prof:
        dispatch_tokens, residual, routing = layer.forward_stage1(next_hidden_states, rotary_posemb, cu_seqlens)
    _prof.print_summary()
else:
    dispatch_tokens, residual, routing = layer.forward_stage1(next_hidden_states, rotary_posemb, cu_seqlens)

if _do_saved:
    _prof = _SavedTensorsProfiler(layer.idx, "stage3", _wt_ptrs)
    with _prof:
        moe_outs = layer.forward_stage3(gathered_tokens, routing.expert_idxs if has_experts else None, routing.expand_idx if has_experts else None)
    _prof.print_summary()
else:
    moe_outs = layer.forward_stage3(gathered_tokens, routing.expert_idxs if has_experts else None, routing.expand_idx if has_experts else None)

if _do_saved:
    _prof = _SavedTensorsProfiler(layer.idx, "stage5", _wt_ptrs)
    with _prof:
        hidden_states = layer.forward_stage5(moe_outs, moe_local_idxs, topk_weight, residual)
    _prof.print_summary()
else:
    hidden_states = layer.forward_stage5(moe_outs, moe_local_idxs, topk_weight, residual)

from pithtrain.pipeline.execution import _layer_mem_profile, _lmem
_do = _layer_mem_profile and self.idx in DETAIL_LAYERS