nvidia/test-writing
> Guide for writing tests in torch-harmonics. Use this skill whenever the user asks to add tests, check or extend test coverage, write a new test class, test a new layer or kernel, or add distributed tests. Also use when the user asks about tolerance values, how to compare tensors, how to structure distributed test infrastructure, or how to test gradients. Covers serial tests in test_convolution.py and distributed tests in test_distributed_*.py.
npx skills add https://github.com/NVIDIA/torch-harmonics --skill test-writingexisting test class in tests/ to pick up naming conventions, helper usage,
and parameterized.expand patterns.
check must compare against a deterministic reference (typically fp32 serial or
the torch reference path). Never just check shapes or dtypes alone.
branches (e.g. fused×kpacked), run each one. If it supports fp32/fp16/bf16,
run all three.
The files for serial tests should be named test_<feature_name>.py
and distributed tests test_distributed_<feature_name>.py respectively.
| 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×).
Tensor comparisons — always use compare_tensors. The first argument is a
name string used in failure messages; it is required:
ok = compare_tensors("output", ref, got, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "output")
Python logic checks — use plain assertTrue / assertFalse:
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.
Every test method that calls compare_tensors must accept a verbose=False
parameter and forward it:
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:
verbose = verbose and self.world_rank == 0
Wrap forward (and any backward that needs grad under the same dtype) with the
maybe_autocast helper from testutils:
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.
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)
Test all three of these in a single test method; do not split them across
separate tests unless there is a strong reason:
compare_tensors..backward(ograd), compare inp.grad with compare_tensors..grad with compare_tensors.See distributed section for the all-reduce needed there.
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")
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):
@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):
...
_DIST_CTX = {}
def setUpModule():
setup_module(_DIST_CTX)
def tearDownModule():
teardown_module(_DIST_CTX)
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.
The pattern is always:
compare_tensors on all ranks; use reduce_success so afailure on any rank fails the test.
# 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()
Use the gather_tensor_hw utility (from testutils) to gather distributed
tensors back to the full shape on all ranks:
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).
Always use reduce_success so a test failure on any rank propagates to all:
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")
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:
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:
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.
disable_tf32() called in setUpClassset_seed(N) called at the top of every test method that uses random tensorsverbose=False parameter on every method that calls compare_tensorssetUpModule / tearDownModule wireupreduce_success used for every assertion in distributed testsmaybe_autocast — not manual .to(dtype) on module weightsTake nvidia/test-writing from the repository into ~/.claude/skills for personal
use, or into .claude/skills inside a project.
The agent identifies a skill by the name field in its header. Two skills with the
same name cannot sit side by side — one of them will be ignored.