---
name: nvidia/test-writing
source: https://app.decimal.ai/s/nvidia-test-writing@1/SKILL.md
source_sha256: 04ab1443aba0
---

# Test Writing Guide

## General principles

- **Match existing style exactly.** Before writing any test, read the relevant
  existing test class in `tests/` to pick up naming conventions, helper usage,
  and `parameterized.expand` patterns.
- **Numerical tests against a reference baseline.** Every forward and backward
  check must compare against a deterministic reference (typically fp32 serial or
  the torch reference path). Never just check shapes or dtypes alone.
- **Test every code path for every precision.** If a layer has 4 dispatch
  branches (e.g. fused×kpacked), run each one. If it supports fp32/fp16/bf16,
  run all three.

---

## Naming Conventions

The files for serial tests should be named `test_<feature_name>.py`
and distributed tests `test_distributed_<feature_name>.py` respectively.

---

## Default tolerances

| dtype | atol | rtol |
|-------|------|------|
| fp64 | 1e-6 | 1e-6 |
| fp32 | 1e-6 | 1e-5 |
| fp16 | 1e-3 | 1e-2 |
| bf16 | 1e-2 | 5e-2 |

Parameter gradients accumulate noise over batch × spatial dims, so loosen their
tolerance by a factor of 10–1000× relative to per-element output tolerances
(use judgment based on reduction size). For stacked/chained layers, loosen
further (2–5×).

---

## Comparison helpers

**Tensor comparisons** — always use `compare_tensors`. The first argument is a
name string used in failure messages; it is required:
```python
ok = compare_tensors("output", ref, got, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "output")
```

**Python logic checks** — use plain `assertTrue` / `assertFalse`:
```python
self.assertIsNotNone(conv.psi_kpacked_K_pad)
self.assertIn(conv.psi_kpacked_K_pad, (8, 16))
self.assertEqual(out.dtype, torch.bfloat16)
```

Never use `torch.allclose` or `torch.testing.assert_close` directly — these
don't give the diagnostic output that `compare_tensors` provides.

### verbose flag

Every test method that calls `compare_tensors` must accept a `verbose=False`
parameter and forward it:

```python
def test_my_feature(self, ..., verbose=False):
    ok = compare_tensors("output", ref, got, atol=atol, rtol=rtol, verbose=verbose)
    self.assertTrue(ok)
```

In distributed tests, only rank 0 should print:
```python
verbose = verbose and self.world_rank == 0
```

---

## Autocast / AMP

Wrap forward (and any backward that needs grad under the same dtype) with the
`maybe_autocast` helper from `testutils`:

```python
with maybe_autocast(self.device.type, dtype):
    out = module(inp)
```

For AMP dtypes (float16, bfloat16), keep module parameters and inputs in
`float32` and let autocast downcast inside the forward. Match this pattern
exactly — do not manually cast module weights or inputs before the call.

```python
is_amp = dtype in (torch.float16, torch.bfloat16)
module_dtype = torch.float32 if is_amp else dtype
module = MyLayer(...).to(dtype=module_dtype, device=device)
inp = torch.randn(..., dtype=module_dtype, device=device)
with maybe_autocast(device.type, dtype):
    out = module(inp)
```

---

## What to test in every test method

Test **all three** of these in a single test method; do not split them across
separate tests unless there is a strong reason:

1. **Forward output** — compare against reference with `compare_tensors`.
2. **Input gradient** — run `.backward(ograd)`, compare `inp.grad` with `compare_tensors`.
3. **Weight (parameter) gradients** — compare each parameter's `.grad` with `compare_tensors`.
   See distributed section for the all-reduce needed there.

```python
inp.requires_grad_(True)
with maybe_autocast(device.type, dtype):
    out = module(inp)
ograd = torch.randn_like(out)
out.backward(ograd)

ok = compare_tensors("output",   ref_out,   out,      atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "output")
ok = compare_tensors("inp grad", ref_igrad, inp.grad, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "inp grad")
if module.weight.grad is not None:
    ok = compare_tensors("weight grad", ref_wgrad, module.weight.grad,
                         atol=wgrad_atol, rtol=wgrad_rtol, verbose=verbose)
    self.assertTrue(ok, "weight grad")
```

---

## Serial test structure

```python
class TestMyFeature(unittest.TestCase):
    device = torch.device("cuda")  # or "cpu"

    @classmethod
    def setUpClass(cls):
        disable_tf32()

    def test_my_case(self, ..., verbose=False):
        set_seed(42)
        ...
```

`disable_tf32()` must be called in `setUpClass` for any test class that runs on
CUDA — TF32 reduces fp32 matmul precision and causes spurious failures at the
default tolerances.

For parameterised cases use `@parameterized.expand([...], skip_on_empty=True)`:

```python
@parameterized.expand([
    # [param1, param2, ..., atol, rtol]
    [64, 128, torch.float32, 1e-6, 1e-5],
    [64, 128, torch.bfloat16, 1e-2, 5e-2],
], skip_on_empty=True)
def test_my_feature(self, nlat, nlon, dtype, atol, rtol, verbose=False):
    ...
```

---

## Distributed test structure

### Module-level wireup (required in every distributed test file)

```python
_DIST_CTX = {}

def setUpModule():
    setup_module(_DIST_CTX)

def tearDownModule():
    teardown_module(_DIST_CTX)
```

### Class setup

```python
class TestMyDistributedFeature(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        setup_class_from_context(cls, _DIST_CTX)
        disable_tf32()

    def setUp(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        if dist.is_initialized():
            dist.barrier()

    def tearDown(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
```

`setup_class_from_context` sets `cls.device`, `cls.world_rank`,
`cls.grid_size_h`, `cls.grid_size_w`, `cls.hrank`, `cls.wrank`,
`cls.h_group`, `cls.w_group` on the class.

`disable_tf32()` is called in `setUpClass` for the same reason as serial tests.

### Comparing distributed output against serial baseline

The pattern is always:
1. Build a serial (local) module and a distributed module with identical weights.
2. Run serial forward + backward on the full tensor.
3. Split the input, run distributed forward + backward on the local shard.
4. Gather distributed outputs/grads back to the full tensor.
5. Compare with `compare_tensors` on all ranks; use `reduce_success` so a
   failure on any rank fails the test.

```python
# 1. build modules
conv_local = MyLayer(**args).to(dtype=module_dtype, device=self.device)
conv_dist  = DistributedMyLayer(**args).to(dtype=module_dtype, device=self.device)

# 2. sync weights — every rank must have identical parameters
with torch.no_grad():
    conv_dist.weight.copy_(conv_local.weight)
    if hasattr(conv_dist, "bias") and conv_dist.bias is not None:
        conv_dist.bias.copy_(conv_local.bias)

# 3. serial fwd + bwd
inp_full = torch.randn((B, C, H, W), dtype=module_dtype, device=self.device)
inp_full.requires_grad_(True)
with maybe_autocast(self.device.type, dtype):
    out_full = conv_local(inp_full)
ograd_full = torch.randn_like(out_full)
out_full.backward(ograd_full)
igrad_full = inp_full.grad.clone()

# 4. distributed fwd + bwd on local shard
inp_local = self._split_helper(inp_full.detach().clone())
inp_local.requires_grad_(True)
with maybe_autocast(self.device.type, dtype):
    out_local = conv_dist(inp_local)
ograd_local = self._split_helper(ograd_full)
out_local.backward(ograd_local)
igrad_local = inp_local.grad.clone()
```

### Gather helpers

Use the `gather_tensor_hw` utility (from `testutils`) to gather distributed
tensors back to the full shape on all ranks:

```python
def _split_helper(self, tensor):
    return split_tensor_hw(tensor,
        hdim=-2, wdim=-1,
        hsize=self.grid_size_h, wsize=self.grid_size_w,
        hrank=self.hrank, wrank=self.wrank)

def _gather_helper_fwd(self, tensor, conv_dist):
    return gather_tensor_hw(tensor,
        hdim=-2, wdim=-1,
        hshapes=conv_dist.lat_out_shapes, wshapes=conv_dist.lon_out_shapes,
        hsize=self.grid_size_h, wsize=self.grid_size_w,
        hrank=self.hrank, wrank=self.wrank,
        hgroup=self.h_group, wgroup=self.w_group)

def _gather_helper_bwd(self, tensor, conv_dist):
    return gather_tensor_hw(tensor,
        hdim=-2, wdim=-1,
        hshapes=conv_dist.lat_in_shapes, wshapes=conv_dist.lon_in_shapes,
        hsize=self.grid_size_h, wsize=self.grid_size_w,
        hrank=self.hrank, wrank=self.wrank,
        hgroup=self.h_group, wgroup=self.w_group)
```

Use `_gather_helper_fwd` for output tensors (indexed by output spatial shape)
and `_gather_helper_bwd` for input-space tensors (input gradient).

### Comparing gathered results

Always use `reduce_success` so a test failure on any rank propagates to all:

```python
verbose = verbose and self.world_rank == 0

out_gather = self._gather_helper_fwd(out_local, conv_dist)
ok = compare_tensors("output", out_full, out_gather, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(reduce_success(ok, self.device), "output")

igrad_gather = self._gather_helper_bwd(igrad_local, conv_dist)
ok = compare_tensors("gradients", igrad_full, igrad_gather, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(reduce_success(ok, self.device), "gradients")
```

### Weight gradient reductions for distributed layers

Distributed layers compute parameter gradients as partial sums over the local
spatial tile. To compare against the serial reference, all-reduce across the
polar (h) and azimuth (w) groups independently:

```python
def _allreduce_param_grad(self, tensor):
    out = tensor.clone()
    if self.grid_size_h > 1:
        dist.all_reduce(out, group=self.h_group)
    if self.grid_size_w > 1:
        dist.all_reduce(out, group=self.w_group)
    return out
```

Then compare:
```python
param_tol = 1000.0 if not is_amp else 10.0   # param grads accumulate more noise
pg_atol, pg_rtol = atol * param_tol, rtol * param_tol
if conv_dist.weight.grad is not None:
    wgrad = self._allreduce_param_grad(conv_dist.weight.grad)
    ok = compare_tensors("weight grad", conv_local.weight.grad, wgrad,
                         atol=pg_atol, rtol=pg_rtol, verbose=verbose)
    self.assertTrue(reduce_success(ok, self.device), "weight grad")
```

The tighter tolerance factor for AMP is intentional: the base atol/rtol for
fp16/bf16 is already 3–4 orders of magnitude looser than fp32, so multiplying
by 1000× makes the bound meaninglessly large.

---

## Checklist before committing a new test

- [ ] `disable_tf32()` called in `setUpClass`
- [ ] `set_seed(N)` called at the top of every test method that uses random tensors
- [ ] `verbose=False` parameter on every method that calls `compare_tensors`
- [ ] All three of forward output, input gradient, weight gradients are tested
- [ ] All relevant precisions are covered (fp32, fp16, bf16 at minimum for CUDA layers)
- [ ] All dispatch branches are covered (e.g. fused=True/False, kpacked-enabled/disabled)
- [ ] Distributed tests use `setUpModule` / `tearDownModule` wireup
- [ ] Distributed weight grads are all-reduced before comparison
- [ ] `reduce_success` used for every assertion in distributed tests
- [ ] AMP tests use `maybe_autocast` — not manual `.to(dtype)` on module weights