Pytorch Integration Guide

Install

pip install --upgrade denpex

Wrap a Training Script

import denpex
import torch
from denpex.sdc import SdcRunner
from denpex.instrument import LayerInstrument

instrument = LayerInstrument(sample_every=10)
instrument.attach(model)

sdc = SdcRunner(canary_inputs=fixed_inputs, every_n_steps=20)
sdc.attach(model)

with denpex.wrap(name="my-gpt-run", tenant="acme"):
    for step, batch in enumerate(dataloader):
        loss = model(batch).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        denpex.report_step(loss.item())

        instrument.flush_if_due(step)
        sdc.run_if_due(step)

Per-layer Metrics

LayerInstrument records on a sampled cadence:

• weight_l2, weight_linf
• grad_l2, grad_linf
• activation_l2
• nan_count, inf_count
• Layer checksum (XXH3)

For FSDP, instrument only on rank 0 (summon_full_params).

SDC Canary

SdcRunner.run_if_due(step) runs a small canary microbatch and compares outputs across ranks. Output bit-for-bit via torch.allclose(atol=1e-7, rtol=1e-5).

DDP Context

denpex.ddp.DdpContext captures:

• is_elastic, restart_count
• host_to_rank_map, group_size
• comm_hook_kind, num_iteration

FSDP Integration

from denpex.fsdp import FsdpInstrument
FsdpInstrument(reshard_after_forward=True).attach(model)

On Crash

Denpex captures stdout/stderr (last 10 MB), wraps the traceback, and emits a DiagnosisEvent whose rule is one of:

• TORCH_CUDA_OOM
• DDP_INIT_FAIL
• GRAD_EXPLOSION
• WEIGHT_EXPLOSION
• NCCL_ASYNC_ERROR
• NCCL_WATCHDOG

What We Do NOT Do

• Denpex does not modify your model.
• Denpex does not read your data tensors.
• Denpex does not export model weights.

Compatibility Tested

• PyTorch 2.0, 2.1, 2.2, 2.3.
• FSDP 2.0+.
• CUDA 12.4+.