HealDA v2 Architecture

[aayushg55]

PhysicsNeMo Pull Request

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR introduces the HealDAv2 v2 architecture for video-based weather data assimilation. It adds a full transformer pipeline: a FiLM observation tokenizer, ragged pixel cross-attention, factorized temporal attention, adaLN-Zero modulation, and context-parallel resharding — all backed by hand-written Triton kernels with autograd custom_op wrappers. Two existing files are updated: HEALPixPatchTokenizer gains a separate_time_axis flag and DiT guards NATTEN latent-grid computation.

• New model stack (healda_v2.py → video_dit.py → video_dit_block.py): HealDAv2 composes a FiLM tokenizer, a per-block PixelCrossAttention, optional temporal attention, and context-parallel resharding inside VideoDiTBlock; the backbone uses adaLN-Zero with zero-initialization for training stability.
• Triton kernels (_obs_tokenizer_kernels.py, _pixel_attn_kernels.py): fused forward/backward kernels for the FiLM tokenizer and ragged GQA pixel cross-attention, registered as torch.library.custom_op with fake-tensor implementations for torch.compile compatibility.
• Autotune cache (triton_autotune_cache.py, pixel_cross_attention.py): per-GPU/rank JSON cache that persists Triton autotune configs across runs; the cache-invalidation hash needs to cover _pixel_attn_kernels.py (not just pixel_cross_attention.py) to catch kernel edits.

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/healda/healda_v2.py	Top-level HealDAv2 model; HealDAv2MetaData sets the deprecated name field, triggering a DeprecationWarning on every instantiation.
physicsnemo/experimental/models/healda/pixel_cross_attention.py	Triton ragged GQA attention + autotune cache; autotune hash only covers this file, not _pixel_attn_kernels.py, so kernel changes silently reuse stale configs; dKV is passed to two kernel argument slots without in-code documentation.
physicsnemo/experimental/models/healda/video_dit.py	VideoDiT backbone; MetaData name is too generic and a blank line before class VideoDiT is missing per PEP 8.
physicsnemo/experimental/models/healda/obs_tokenizer.py	FiLM observation tokenizer with Triton fused kernel + autograd wrappers; logic is sound, custom-op/fake registration follows correct patterns.
physicsnemo/experimental/models/healda/pixel_attention_utils.py	Ragged packing utilities (counting sort, cu_seqlens, group map); int32 position cast in counting_sort_and_pack silently overflows for N > ~2 billion observations.
physicsnemo/experimental/models/healda/video_dit_block.py	DiT block with spatial/temporal/cross-attention; adaLN-Zero wiring, drop-path, and context-parallel resharding look correct.
physicsnemo/experimental/models/healda/obs_context.py	ObsContext/PixelGroupMap dataclasses with device-transfer helpers; straightforward and correct.
physicsnemo/experimental/models/healda/adaln.py	adaLN-Zero modulation, ndim-agnostic broadcast; logic is correct for both n_blocks=1 and n_blocks=2 cases.
physicsnemo/nn/module/hpx/tokenizer.py	HEALPixPatchTokenizer gains separate_time_axis flag; detokenizer handles both 3D and 4D inputs; CalendarEmbedding error message improved. Changes are clean.
physicsnemo/models/dit/dit.py	DiT guards latent-grid computation behind is_natten; safe and minimal change.
physicsnemo/experimental/models/healda/temporal_attention.py	Temporal self-attention with RoPE, linear/softmax variants, and causal windowing; @torch.compile on forward looks correct.
physicsnemo/experimental/models/healda/sharding.py	Context-parallel t<->x resharding via all_to_all and ShardTensor; both paths are autograd-aware and logically consistent.

Comments Outside Diff (3)

1. physicsnemo/experimental/models/healda/pixel_cross_attention.py, line 1465-1481 (link)

   Cache hash reads pixel_cross_attention.py, not the actual Triton kernel file

   The cache-busting digest is computed over __file__ (i.e., pixel_cross_attention.py), but the kernels that actually get tuned live in _pixel_attn_kernels.py. A change to those kernels without touching this file will silently reuse stale autotune configs that were optimized for the old kernel, potentially applying wrong launch parameters to new code.

2. physicsnemo/experimental/models/healda/video_dit.py, line 1349-1365 (link)

   Generic class name MetaData risks import collisions and lacks a blank line

   MetaData is a very common identifier — any from physicsnemo.experimental.models.healda.video_dit import MetaData (or star import) will shadow other MetaData symbols in the same namespace. A more specific name like VideoDiTMetaData avoids this. Also, PEP 8 requires two blank lines before a top-level class definition, but there is only one blank line separating MetaData from class VideoDiT.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

3. physicsnemo/experimental/models/healda/pixel_attention_utils.py, line 273-307 (link)

   counting_sort_and_pack: atomic-claimed position cast to int32 can overflow for very large observation counts

   The Triton kernel computes pos = tl.atomic_add(bucket_offsets_ptr + key, 1, ...) and then writes sorted_order_ptr + pos.to(tl.int32). If the total number of observations N exceeds ~2.1 billion, the int32 cast truncates the position, writing two source indices to the same slot and silently corrupting the permutation. The sorted_order output tensor is also torch.int32. For typical atmospheric observation batch sizes this is unlikely, but it is an undocumented limit worth noting in the docstring.

_{Reviews (1): Last reviewed commit: "test(experimental/healda): rename crypti..." | Re-trigger Greptile}

[greptile-apps]

Deprecated name field triggers DeprecationWarning on every instantiation

ModelMetaData.__post_init__ fires a DeprecationWarning whenever name is not _DEPRECATED_SENTINEL. Because HealDAv2MetaData sets name: str = "HealDAv2", every HealDAv2(...) call emits this warning. The name field was removed from the base class and is replaced by the model's class name automatically — the field should be dropped from HealDAv2MetaData entirely.

Suggested change

	cuda_graphs: bool = False
	amp_cpu: bool = False
	amp_gpu: bool = True
	torch_fx: bool = False
	bf16: bool = True
	onnx: bool = False
	func_torch: bool = False
	auto_grad: bool = False
	class HealDAv2(Module):
	r"""
	Video data-assimilation model combining HEALPix tokenizers, a FiLM observation
	tokenizer, and a VideoDiT backbone.
	@dataclass
	class HealDAv2MetaData(ModelMetaData):
	"""Metadata for HealDAv2 model."""
	jit: bool = False

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Same dKV tensor passed to two separate kernel pointer slots

dKV appears at two consecutive positions in the _pixel_attn_gqa_bwd call. With COMBINED_DKV=True the kernel presumably writes the combined [dK | dV] rows to only the first argument and ignores the second, but this aliased pointer pattern is not documented in this call site. If a future kernel change removes the COMBINED_DKV guard and begins writing to both pointers independently, the two writes will race on the same backing storage, silently corrupting gradients. At a minimum, add a comment here explaining why the same tensor is safe to pass twice.

[aayushg55]

@pzharrington Let me know what you think about the VideoDiTBlock/VideoDiT classes. I can restructure the code accordingly.

In particular, would something like this work better?

class VideoDiTBlock(DiTBlock):
    def __init__(self, ..., temporal_attention=False, cross_attention=None, is_causal=False, ...):
        super().__init__(hidden_size, num_heads, attention_backend=..., condition_embed_dim=..., mlp_ratio=..., drop_path=..., **attn_kwargs)
        self._is_causal = is_causal
        if temporal_attention: self.temporal_attention = ...; self.temporal_attn_norm = ...
        if cross_attention:    self.cross_attention = cross_attention(); self.cross_attn_norm = ...
    def forward(self, hidden_states, c, cross_attention_context=None, attn_kwargs=None):
        # reuse self.adaptive_modulation(c).chunk(6), self.pre_attention_norm/self.pre_mlp_norm,
        # self.attention, self.linear, self.drop_path; insert obs + temporal sub-layers between.

[aayushg55]

The video block has four gated sub-layers (spatial attention, MLP, temporal attention, obs cross-attention). So rather than repeat DiTBlock's inline adaLN pattern — adaptive_modulation = Sequential(SiLU, Linear(cond, 6*hidden)) + separate pre_attention_norm/pre_mlp_norm + a modulation helper — four times, I felt it would be cleaner to wrap the whole adaLN-Zero operation (SiLU + Linear → shift/scale/gate, the affine-free LayerNorm, and applying them) into a single reusable AdaLayerNormZero module.

[aayushg55]

Context parallelism: manual all-to-all vs ShardTensor.redistribute

Factorized temporal attention needs a time↔space reshard each block (2 per block:
enter t-sharded → spatial attn → reshard to x-sharded → temporal attn → reshard
back). We implemented it both ways: a manual all_to_all_single, and
ShardTensor.redistribute (Shard(time) ↔ Shard(space) on a 1D mesh).

Per-block — 8×H100, dit-5B (C=1536 (16×96), X=12288, T=8, bf16 autocast, 2 reshards/block:

path	eager ms/block	compiled ms/block
manual all-to-all	16.5	14.5
ShardTensor.redistribute	20.5	19.3

Just the reshard primitives (t↔x round trip = 2 reshards, same shape, eager):

path	fwd	fwd+bwd
manual all-to-all	0.29 ms	0.97 ms
ShardTensor.redistribute	2.48 ms	5.45 ms

[aayushg55]

Benchmark code — reshard primitives (reshard_microbench.py)

Run on 8×H100, dit-5B shape:
T_FULL=8 X_FULL=12288 EMBED_DIM=1536 torchrun --nproc_per_node=8 reshard_microbench.py

import os
import time

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.placement_types import Shard

from physicsnemo.domain_parallel import ShardTensor
from physicsnemo.experimental.models.healda.sharding import (
    shard_t,
    shard_t_shardtensor,
    shard_x,
    shard_x_shardtensor,
)

dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()
torch.cuda.set_device(rank)
dev = torch.device("cuda", rank)
mesh = init_device_mesh("cuda", (world,), mesh_dim_names=["cp"])
group = mesh.get_group()

b = 1
t_full = int(os.environ.get("T_FULL", 16))
x_full = int(os.environ.get("X_FULL", 12288))
c = int(os.environ.get("EMBED_DIM", 1536))
t_local = t_full // world
dtype = torch.bfloat16


def manual(local):
    return shard_t(shard_x(local, group), group)


def st(local):
    return shard_t_shardtensor(shard_x_shardtensor(local, mesh), mesh)


def bench(fn, name, backward, iters=30, warmup=10):
    local = torch.randn(
        b, t_local, x_full, c, device=dev, dtype=dtype, requires_grad=backward
    )
    for _ in range(warmup):
        out = fn(local)
        if backward:
            out.float().pow(2).mean().backward()
            local.grad = None
    torch.cuda.synchronize()
    dist.barrier()
    t0 = time.perf_counter()
    for _ in range(iters):
        out = fn(local)
        if backward:
            out.float().pow(2).mean().backward()
            local.grad = None
    torch.cuda.synchronize()
    dist.barrier()
    ms = (time.perf_counter() - t0) / iters * 1e3
    if rank == 0:
        print(f"  {name:12s} {'fwd+bwd' if backward else 'fwd':8s}: {ms:8.3f} ms/iter")


def bench_stage(name, fn, iters=30, warmup=10):
    local = torch.randn(b, t_local, x_full, c, device=dev, dtype=dtype)
    for _ in range(warmup):
        fn(local)
    torch.cuda.synchronize()
    dist.barrier()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn(local)
    torch.cuda.synchronize()
    dist.barrier()
    if rank == 0:
        print(f"    stage {name:24s}: {(time.perf_counter() - t0) / iters * 1e3:8.3f} ms")


if rank == 0:
    print(f"world={world} local=({b},{t_local},{x_full},{c}) {dtype}")
for bwd in (False, True):
    bench(manual, "manual a2a", bwd)
    bench(st, "shardtensor", bwd)

_st = ShardTensor.from_local(
    torch.randn(b, t_local, x_full, c, device=dev, dtype=dtype), mesh, [Shard(1)]
)
bench_stage("from_local", lambda x: ShardTensor.from_local(x, mesh, [Shard(1)]))
bench_stage("redistribute", lambda x: _st.redistribute(placements=[Shard(2)]))
dist.destroy_process_group()

[aayushg55]

@negin513 for viz

[pzharrington]

I suspect the ShardTensor dispatch overhead is responsible to the differences in eager mode. ShardTensor does not yet support torch compilation, but will very very soon (like a week or two). Could be interesting to look at this again when it lands. It'd be nice to have a fully ShardTensor-based implementation but if this PR is looking ready before then, perhaps we can merge it with both implementations and defer the decision (which should not change the user-facing APIs/functionality

[aayushg55]

Our model needs QK-normalization with affine=False in the spatial attention, but TE's MultiheadAttention only supports affine=True. Options I see:

• (a) TE: Keep the affine params but exclude them from the optimizer.
• (b) Build a small attention module: affine-free QK-norm + TE DotProductAttention.
• (c) Use the existing timm backend, which already supports affine=False.

Currently, I went with c (this is also what the v1 ckpt uses), but just wanted to flag this incompatibility with the TE backend as set up now.

[aayushg55]

This is a patch to fix a recent DiT / HealDA v1 bug (unrelated to my PR introducing v2)