GLOBE Performance Improvements

[peterdsharpe]

PhysicsNeMo Pull Request

This PR is a focused performance pass on the experimental GLOBE model and its training scripts. No model architecture changes, no checkpoint compat changes; all numerics are unchanged at fp32-epsilon precision (a couple of internal aggregations now happen in fp64 - details in §5).

Net results on 16x B200s, DrivAerML training with standalone recipe: per-sample time (load, forward, backward, step) goes from ~7.8 seconds to ~5.5 seconds. (-30%)

TL;DR

• CPU↔GPU synchronization is the dominant win. The hot path used to issue ~60+ implicit syncs per training step (mem_get_info, .item(), if X.any():, tensor[bool_mask], …). Most have been eliminated.
• Two O(tree_depth) Python loops are gone. Per-node source aggregation and per-node strength accumulation are now single cumsum + range-subtract ops, exploiting the morton-sorted source ordering.
• The four-phase BarnesHutKernel scatter is packed. One index_add_ per phase instead of one scatter_add_ per phase per output field. This more than halves the backward-pass scatter cost on DrivAerML, where indexing_backward was ~23 % of total GPU time per step.
• GradScaler is gone from the training scripts (it's a no-op with bf16 autocast and forces a per-step inf/NaN-check sync).
• CUDA allocator is now expandable_segments:True to avoid synchronizing cudaMalloc / cudaFree round-trips under the chunked kernel workload.
• Numerical-stability fix: centroid / node-strength aggregations now cumsum in fp64. Was producing p99 ~100 % wrong leaf centroids at DrivAer scale (N=1M, coords ~5 m, leaf_size=1).
• Several torch.compile graph breaks are fixed, mostly as a side-effect of the changes above.
• New optional tree_build_device knob ("cpu" / "cuda" / None) for the small-problem regime where CUDA-launch latency dominates tree work.

Why these changes matter (hot path context)

A GLOBE training step is dominated by two things:

1. Per-sample geometry preprocessing. Build one ClusterTree per BC type, then run a dual-tree Barnes-Hut traversal for every (src BC, dst BC) pair to produce a DualInteractionPlan describing which target/source point pairs interact in the near, far, and mixed-rank phases.
2. BarnesHutKernel evaluation. Apply MLPs to those pairs and scatter_add the results back to per-target buffers, four times (one for each near/far phase), once per multiscale branch, once per communication hyperlayer.

Profiling showed both phases were CPU-bound on launch/sync overhead rather than GPU-bound on the actual MLP and indexing work. The changes below attack that bottleneck.

---

1. Eliminating GPU↔CPU synchronization

Each entry below removes a sync point that profiling identified as significant. These are the same change in spirit (avoid materializing int(tensor) or branching on tensor contents in the hot path), applied to several different call sites.

1.1 _auto_chunk_size: mem_get_info → cached total_memory

BarnesHutKernel._auto_chunk_size ran on every kernel evaluation and called torch.cuda.mem_get_info(device), which is a synchronizing driver query. Across the four phases × multiple branches × multiple hyperlayers, this was ~60 syncs per training step. Replaced with a cached lookup of total device memory and a fixed-fraction budget (25 %):

_CHUNK_MEMORY_BUDGET_PERCENT: Final[int] = 25


def _ceil_div(a: int, b: int) -> int:
    # ... pure-integer ceiling division ...
    return -(-a // b)


@torch.compiler.disable
@cache
def _device_total_memory_bytes(device: torch.device) -> int:
    # CUDA: torch.cuda.get_device_properties(device).total_memory
    # CPU:  psutil.virtual_memory().total
    ...

@torch.compiler.disable
def _device_chunk_budget_bytes(device: torch.device) -> int:
    return _device_total_memory_bytes(device) * _CHUNK_MEMORY_BUDGET_PERCENT // 100

Same swap is applied to MultiscaleKernel.forward's branch-checkpoint decision.

The trade is "exactly fits in currently-free memory" → "fits in 25 % of total device memory". Fine on 100+ GB devices (e.g. GB200) where the non-kernel resident state is well under the remaining 75 %.

1.2 Smoothing radius: Python float → registered buffer

Kernel._evaluate_interactions was rebuilding a tensor on every forward pass:

smoothing_radius = torch.tensor(self.smoothing_radius, device=device, dtype=dtype)
... (vectors * vectors).sum(dim=-1).apply(lambda x: x + smoothing_radius**2)

Now stored as a registered buffer and used directly with TensorDict broadcasting:

        ### Pre-squared smoothing radius as a registered tensor buffer.
        ### A buffer (rather than a Python float) is required because
        ### ``_evaluate_interactions`` adds this scalar to a TensorDict
        ### inside a checkpoint sub-graph; Python free variables cannot
        ### be lifted across nested Dynamo SubgraphTracers (the lift
        ### chain bottoms out at the root tracer with no parent and
        ### asserts ``lift_tracked_freevar_to_input should not be
        ### called on root SubgraphTracer``).  As a buffer, the
        ### tensor is a tracked module attribute that Dynamo treats as
        ### a graph leaf, so no lift is needed.
        self.register_buffer(
            "_smoothing_radius_sq",
            torch.tensor(smoothing_radius**2, dtype=torch.float32),
            persistent=False,
        )

Eager: skips a tensor allocation and a .apply(lambda) per forward. Compile: closes the torch.compile graph break.

1.3 Cluster tree construction: per-level sync → end-of-loop compaction

ClusterTree.from_points used to do, on every level of every tree,

leaf_indices = torch.where(is_leaf_seg)[0]
if len(leaf_indices) > 0:
    ...

That's a CPU-GPU sync per level × per tree (the in-code comment notes ~16 levels × 4 trees × 28 samples for DrivAerML). The loop now accumulates (seg_node_ids, seg_starts, seg_sizes, is_leaf_seg) per level and pays one boolean compaction after the loop:

            ### Defer leaf-segment processing to a single end-of-loop pass.
            # Per-iter ``torch.where(is_leaf_seg)[0]`` was a CPU-GPU sync
            # point on every level of every tree (~16 levels x 4 trees x
            # 28 samples).  We instead accumulate (node_id, start, size,
            # validity) for each segment seen during the loop and pay one
            # boolean compaction at the end.  ``torch.where(is_internal_seg)``
            # remains in-loop because the next iteration's active segments
            # are derived from this iteration's internals.
            leaf_seg_node_ids: list[torch.Tensor] = []
            leaf_seg_starts: list[torch.Tensor] = []
            leaf_seg_sizes: list[torch.Tensor] = []
            leaf_seg_validity: list[torch.Tensor] = []

The previous helpers _fill_leaf_aabbs and _fill_leaf_total_areas also each called int(leaf_sizes.sum()); they're merged into a single _fill_leaf_aggregates that shares the _ragged_arange mapping between AABB and area aggregation, removing the redundant .sum() sync.

1.4 Dual-tree traversal: ~12 .any()-gated indexings → one compaction

ClusterTree.find_dual_interaction_pairs is the hot loop of geometry prep. The previous body did 12 separate if X.any(): early-exits and tensor[bool_mask] selections per iteration - each one a sync.

The new body classifies every active node-pair into one of three cases (case_T_only, case_S_only, case_both), computes all eight potential child slots with validity masks, stacks them, and pays one boolean compaction per iteration:

                ### 3. Generate next iteration's active set.
                # We compute children over the FULL active set (n_active
                # entries) and use validity masks per (T,S) child slot to
                # encode the case-A / case-B / case-C splitting rules from
                # the original implementation.  After unioning the eight
                # potential child slots we pay ONE boolean compaction
                # instead of the original ~12 ``.any()``-gated indexings.
                ...
                slot_t = torch.stack([...])
                slot_s = torch.stack([...])
                slot_v = torch.stack([...])
                ...
                flat_v = slot_v.reshape(-1)
                active_tgt_nodes = slot_t.reshape(-1)[flat_v]
                active_src_nodes = slot_s.reshape(-1)[flat_v]

Far-field outputs are similarly accumulated unfiltered into per-iteration lists with a parallel is_far validity mask, then compacted ONCE at the very end of the traversal.

1.5 Leaf-pair expansion: dedup nonzero calls

In _expand_dual_leaf_hits, tensor[bool_mask] lowers to a synchronizing aten::nonzero that sizes the output. When several indexings share the same mask, the previous code paid one sync per indexing. They're now deduped:

    ### (near, far) output.  ``target_is_far`` is consumed by two
    ### indexings; doing one ``nonzero`` and reusing the integer index
    ### saves one sync (each ``tensor[bool_mask]`` lowers to a
    ### synchronizing ``aten::nonzero`` to size the output).
    far_idx_t = target_is_far.nonzero(as_tuple=True)[0]
    nf_target_ids = target_point_ids[far_idx_t]
    nf_source_node_ids = src_leaf_per_target[far_idx_t]

The function also drops a pile of "early return on empty input" branches that were each a sync; the downstream ops (scatter, bincount, repeat_interleave, _ragged_arange) all handle zero-element inputs correctly.

1.6 Traversal loop bound: max_depth.item() → bit_length

The dual-traversal loop bound was

max_iters = int(target_tree.max_depth.item()) + int(source_tree.max_depth.item()) + 1

i.e. two .item() calls before the loop starts. Replaced with a static upper bound that needs no GPU read:

            ### Loop bound: every iteration descends at least one tree level
            ### on at least one side, so ``2 * total_levels + safety`` is a
            ### hard upper bound that requires no GPU->CPU read.  Using
            ### ``int(max_depth.item())`` as before would force two syncs
            ### per call before we even start the loop.
            n_src_levels = max(1, int(source_tree.n_sources).bit_length())
            n_tgt_levels = max(1, int(target_tree.n_sources).bit_length())
            max_iters = 2 * (n_src_levels + n_tgt_levels) + 4

n_sources is a Python int (a .shape[0] lookup), not a tensor.

1.7 Tree-build wrapper: no_grad + minor cleanups

GLOBE._build_trees_and_plans and _build_prediction_plans now run inside torch.no_grad() (the outputs are integer indices and a non-grad-tracked area divisor, so autograd bookkeeping on dozens of tensor ops was wasted) and inline the per-plan debug logging into the build loop instead of re-iterating afterward.

---

2. Prefix-sum aggregation: replacing Python loops over tree levels

A cluster tree's morton-sorted source ordering has a key property: every node covers a contiguous range of sources, [node_range_start, node_range_start + node_range_count). So any node-subtree sum is a prefix-sum range subtract: prefix[end] - prefix[start].

This kills two O(tree_depth) Python loops that were previously the dominant cost in their respective functions.

2.1 ClusterTree.compute_source_aggregates

The previous implementation did:

• Per-leaf weighted-sum via segmented scatter.
• A Python-level for level_ids in reversed(depth_levels): propagating centroids and per-source features bottom-up.

Combined cost in profiling: ~2 s per training step of mixed CPU (level-loop overhead) and GPU (kernel launches per level).

Replaced with a single cumsum + range-subtract over the morton-sorted sources:

        ### Range-sum aggregation via morton-sorted prefix sums.
        # Each node covers a contiguous range
        # [node_range_start, node_range_start + node_range_count) in
        # morton-sorted source order, so any node-subtree sum is just
        # ``prefix[end] - prefix[start]``.  This replaces the old
        # leaf-aggregation + bottom-up Python loop, which were the
        # dominant CPU + GPU costs in ``compute_source_aggregates``
        # (~2 s combined per training step in profiling).
        ...
        sorted_points = source_points[self.sorted_source_order]
        sorted_areas = areas[self.sorted_source_order]
        weighted_points_64 = (sorted_points * sorted_areas.unsqueeze(-1)).double()
        cumsum_weighted_points = torch.nn.functional.pad(
            torch.cumsum(weighted_points_64, dim=0), (0, 0, 1, 0)
        )
        ...
        node_total_weighted_pts = (
            cumsum_weighted_points[ends] - cumsum_weighted_points[starts]
        )

Per-source feature aggregation gets the same treatment via an inner _aggregate_via_prefix_sum, applied across TensorDict leaves. Trailing feature dims are flattened so cumsum sees a single feature axis (avoiding a per-feature kernel chain inside cumsum).

The internal_level_ids / internal_level_offsets / leaf_node_ids / leaf_seg_ids tensors stored on ClusterTree purely to support the old bottom-up loop are dropped.

2.2 BarnesHutKernel._compute_node_strengths

Same story, on a single scalar field (source_strengths). The old implementation did leaf-summing via scatter then a bottom-up Python loop; the new one is one cumsum + range-subtract:

        ### Cumsum and range-subtract in fp64 to avoid catastrophic
        ### cancellation when ``cumsum_total >> range_sum`` - the regime
        ### of small leaves in a large tree built over offset coordinates.
        ### See the matching note in :meth:`ClusterTree.compute_source_aggregates`.
        sorted_strengths_64 = source_strengths[tree.sorted_source_order].double()
        ### Pad with a leading zero so that ``cumsum[i]`` is the sum of
        ### sorted_strengths[:i] - both endpoints index identically.
        prefix_sum = torch.nn.functional.pad(
            torch.cumsum(sorted_strengths_64, dim=0), (1, 0)
        )

        starts = tree.node_range_start
        ends = starts + tree.node_range_count
        return (prefix_sum[ends] - prefix_sum[starts]).to(source_strengths.dtype)

---

3. Packed scatter for backward-pass efficiency

BarnesHutKernel scatters into per-target output buffers four times per forward (near, far-node, near-far, far-near), once per output field. In profiling, indexing_backward_kernel_stride_1 (the backward of scatter_add_) was ~1 s out of ~4.5 s total GPU time per training step (~23 %) - the largest single GPU kernel by time.

Two changes in BarnesHutKernel.forward:

1. Pack output fields into a single (n_targets, total_features) buffer instead of per-field buffers.
2. Use index_add_ instead of scatter_add_ for the per-target accumulation.

        ### Packed output buffer.  All four phases scatter into this single
        ### tensor; the per-phase loops over output keys (one ``scatter_add_``
        ### per key) are replaced with one packed ``scatter_add_`` per phase.
        ### ``indexing_backward`` was the top GPU kernel by time in profiling
        ### (~23% of total GPU time); packing cuts this by the number of
        ### output fields (~4x for ``C_p`` + ``C_f`` on DrivAerML).
        ...
        packed_buf = torch.zeros(
            (n_targets, total_features), dtype=buffer_dtype, device=device
        )

The pack/scatter helper is shared across all four phases:

    def _pack_and_scatter(
        self,
        chunk_result: TensorDict,
        weights: Float[torch.Tensor, " n_pairs"],
        tgt_ids: Int[torch.Tensor, " n_pairs_or_expanded"],
        packed_buf: Float[torch.Tensor, "n_targets total_features"],
        *,
        broadcast_pair_ids: Int[torch.Tensor, " n_pairs_or_expanded"] | None = None,
    ) -> None:
        ...
        ### ``index_add_`` rather than ``scatter_add_`` with broadcasted
        ### indices: equivalent semantics, but ``index_add_`` takes a 1-D
        ### ``index`` and avoids the ``unsqueeze`` + ``expand_as`` overhead
        ### that ``scatter_add_`` requires.  In addition to the direct
        ### speedup, ``index_add_`` has a more compact backward (no index
        ### broadcasting in the saved tensors), which compounds with the
        ### packing optimization above.
        packed_buf.index_add_(0, tgt_ids, weighted)

Counterpart _unpack_buf slices back into a per-field TensorDict. The canonical-ordered field list and per-key feature widths are cached as Kernel._output_packing.

Net effect on DrivAerML (C_p 1ch + C_f 3ch = 4 fields, 4 phases): 16 scatter_add_ calls → 4 index_add_ calls per kernel evaluation.

---

4. Memory management

4.1 CUDA allocator: expandable_segments:True

Both run.sh scripts now export

export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

The chunked kernel evaluations stress the cache enough that the default segment allocator falls back to synchronous cudaMalloc / cudaFree round-trips. expandable_segments lets the allocator grow existing segments instead.

4.2 Gradient checkpointing applies in both train and eval

The four if self.training and self.use_gradient_checkpointing branches across BarnesHutKernel.forward are hoisted into one helper:

    def _maybe_checkpointed_evaluate(self, *args: object) -> TensorDict:
        """Run :meth:`_gather_and_evaluate` with optional gradient checkpointing.

        ...

        The wrap fires in **both** training and eval, not just training,
        for two reasons:

        1. Memory savings (the original purpose).  In eval no autograd
           tape is active (``train.py`` wraps validation in
           ``torch.no_grad()``), so ``checkpoint(use_reentrant=False)``
           degenerates to a near-no-op forward call - no recompute, no
           extra memory, no extra compute.
        2. Workaround for a Dynamo+CUDA TensorDict bug.  When the same
           ``_gather_and_evaluate`` body is inlined into the parent
           graph (i.e. without the checkpoint wrapper), FakeTensor
           tracing fails to propagate ``TensorDict.batch_size`` through
           ``(vectors * vectors).sum(dim=-1)`` ...
        """
        if self.use_gradient_checkpointing:
            return checkpoint(
                self._gather_and_evaluate, *args, use_reentrant=False
            )
        return self._gather_and_evaluate(*args)

The self.use_gradient_checkpointing flag is also piped through MultiscaleKernel → Kernel → BarnesHutKernel and exposed as a constructor kwarg on GLOBE (default True).

4.3 Packed output buffer

BarnesHutKernel.forward used to lazily allocate per-field, per-phase output buffers (output_bufs[k] = torch.zeros(...) if-key-not-seen pattern). Now allocated up-front:

        ### Buffer dtype must match the dtype of ``weighted = chunk * weights``
        ### that the four phases will scatter in (``index_add_`` requires
        ### exact dtype match).  ``chunk`` comes from the MLP at the
        ### autocast dtype (or ``source_points.dtype`` outside autocast);
        ### ``weights`` carries ``source_strengths.dtype``.  Their product
        ### is the type-promoted dtype, which is what the previous
        ### lazy-allocated buffers used to capture - we now compute it
        ### eagerly so a single buffer can be allocated up front.
        ...
        packed_buf = torch.zeros(
            (n_targets, total_features), dtype=buffer_dtype, device=device
        )

Composes with §3 (packed scatter) and lets us drop the if n_near == 0 and n_nf == 0 and n_fn == 0 and n_far_nodes == 0 empty-result short-circuit at the top of forward: an all-zero packed_buf unpacks to all-zero per-field tensors, which is the correct result anyway.

---

5. Numerical stability via fp64 cumsum

The cumsum-based aggregations in §2 are done in float64 and cast back to source_points.dtype at the end. Reason from the code:

        # The cumsum and the range subtract are done in fp64 because fp32
        # suffers catastrophic cancellation when ``range_sum << cumsum_total``,
        # which is the regime of small leaves (``leaf_size=1``) in a large
        # tree built over offset (e.g. all-positive) coordinates.  At
        # drivaer scale (``N=1M``, coords ~5 m), fp32 leaf centroids had
        # median ~2 % relative error and p99 ~100 % wrong.  Lifting the
        # cumsum to fp64 brings this back to fp32 epsilon (~1e-7) and adds
        # <1 % wall-clock to the training step (cumsum is ~2.3x slower in
        # fp64, but cumsum is a tiny fraction of step time).  CUDA fp32
        # cumsum is also non-deterministic across runs (pytorch#75240);
        # fp64 cumsum is much less affected.

Listed under "performance" because it's a prerequisite for §2 - without it the prefix-sum trick can't be used.

---

6. torch.compile compatibility

Most of the syncs above (§1) and the packed-scatter rewrite (§3) also fix torch.compile graph breaks. A few changes here are specifically for torch.compile:

• _ceil_div for chunk sizing instead of math.ceil(a / b). Under specialize_float=False, a Python float gets traced as an unbacked symbolic float that propagates into the chunk size and crashes Dynamo at the checkpoint boundary in _gather_and_evaluate. Pure-integer ceiling division (-(-a // b)) sidesteps this.
• _CHUNK_MEMORY_BUDGET_PERCENT: Final[int] = 25 instead of e.g. 0.25 - same reason.
• @torch.compiler.disable on the device-memory helpers and _auto_chunk_size, so Dynamo doesn't try to evaluate torch.cuda.get_device_properties as a graph constant (which crashes on CPU-only hosts and produces unbacked symbolic ints on CUDA hosts).
• The smoothing-radius buffer from §1.2 also fixes the lift-across- SubgraphTracers issue.
• view(-1, 1) → reshape(-1, 1) in GLOBE.forward's output calibration loop, for robustness when upstream ops produce non-contiguous tensors.

The hierarchical_acceleration.md design doc is updated to reflect the new "static fraction of total device memory" chunk-size strategy.

---

7. Training scripts

Both examples/cfd/external_aerodynamics/globe/airfrans/train.py and .../drivaer/train.py have the same three changes.

7.1 Removed GradScaler

GradScaler is only useful for fp16 autocast (to combat gradient underflow). The training scripts use torch.bfloat16 autocast, which has the same dynamic range as fp32 - there's no underflow to protect against. When enabled=amp, GradScaler still forces a per-step inf/NaN check sync on the gradients, for nothing.

-    scaler = torch.amp.GradScaler(device=device.type, enabled=amp)
-    ...
-    scaler.scale(batch_loss).backward()
+    batch_loss.backward()
     if gradient_clip_norm is not None:
-        scaler.unscale_(optimizer)
         torch.nn.utils.clip_grad_norm_(
             model.parameters(), max_norm=gradient_clip_norm
         )
-    scaler.step(optimizer)
-    scaler.update()
+    optimizer.step()

(Removed from checkpoint save/load plumbing too.)

7.2 New tree_build_device knob

Exposed as a Literal["cpu", "cuda"] | None argument:

        leaf_size=leaf_size,
        network_type=network_type,
        self_regularization_beta=self_regularization_beta,
        latent_compression_scale=latent_compression_scale,
        expand_far_targets=expand_far_targets,
        tree_build_device=tree_build_device,
    ).to(device)

Forwarded to GLOBE.__init__, which transfers boundary centroids/areas to build_device for from_points + find_dual_interaction_pairs, then moves the resulting ClusterTree and DualInteractionPlan (both @tensorclass, so one .to() call each) back to the input's device at the end of _build_trees_and_plans / _build_prediction_plans.

Default None preserves prior behavior (build on the input's device). "cpu" can win for small problems (~a few thousand boundary cells) where CUDA launch latency + cudaStreamSynchronize round-trips dominate the actual tree work.

7.3 CUDA allocator env var (run.sh)

See §4.1.

---

8. Dependencies (pyproject.toml)

-    "tensordict>=0.10.0",
+    "tensordict>=0.12.2",
     "omegaconf>=2.3.0",
     "importlib-metadata>=8.7.1",
+    "psutil>=6.0.0",

• tensordict>=0.12.2 (was >=0.10.0) - bumped alongside the BarnesHutKernel refactor; the codebase exercises TensorDict broadcasting and batch-size propagation more heavily than before (see §4.2 for one such case in _maybe_checkpointed_evaluate), and this update brings better torch.compile support.
• psutil>=6.0.0 - used by _device_total_memory_bytes for the CPU branch (psutil.virtual_memory().total). The CPU build path is debug-only; production runs on CUDA, where the CUDA branch is used.

---

What's intentionally not in this PR

• No model-architecture changes. No new layers, no new losses, no hyperparameter changes.
• No checkpoint format changes - existing .mdlus checkpoints load unmodified.
• The two new constructor kwargs on GLOBE (use_gradient_checkpointing, tree_build_device) both have defaults matching the prior behavior, so this is a non-breaking change for downstream users.

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

A focused performance pass on the GLOBE model achieving ~30% per-step speedup on 16× B200s. The dominant wins are eliminating ~60 GPU↔CPU sync points per training step (deferred boolean compactions, static memory budget, bit_length-based loop bound), replacing two O(tree_depth) Python loops with O(N) prefix-sum range subtracts, and packing the four-phase Barnes-Hut scatter into a single index_add_ per phase. A real numerical bug (fp32 cumsum catastrophic cancellation giving ~100% wrong leaf centroids at DrivAer scale) is also fixed as a prerequisite.

• Sync elimination (cluster_tree.py, field_kernel.py): per-iteration .any() / .item() / boolean-mask indexings replaced with one deferred compaction per traversal loop iteration; mem_get_info per kernel call replaced with a cached 25%-of-total-memory budget.
• Packed scatter (field_kernel.py): all four BH phases now scatter into a single (n_targets, total_features) buffer via index_add_, cutting indexing_backward from ~23% to ~6% of GPU time.
• Training script cleanup (train.py ×2): GradScaler removed (correct no-op for bfloat16 autocast), tree_build_device kwarg added, PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True exported.

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/globe/cluster_tree.py	Major refactor eliminating ~60+ GPU↔CPU sync points: deferred leaf compaction, single-compaction dual-traversal, and O(N) prefix-sum aggregation replacing O(tree_depth) Python loops. Logic verified correct; new fp64 cumsum fixes real centroid precision bug at DrivAer scale.
physicsnemo/experimental/models/globe/field_kernel.py	Packed scatter buffer (index_add_ replacing per-field scatter_add_), smoothing radius registered buffer, and static memory budget. Two P2 concerns: checkpoint now applies in eval mode by default, and fixed 25% memory budget may exceed available VRAM on loaded devices.
physicsnemo/experimental/models/globe/model.py	Adds tree_build_device and use_gradient_checkpointing kwargs, wraps tree-build passes in torch.no_grad(), and removes the overly defensive .exclude("strengths") filter. reference_area is a registered buffer so .device fallback is valid. view→reshape fix is correct.
examples/cfd/external_aerodynamics/globe/airfrans/train.py	GradScaler removed (correct for bfloat16 autocast), tree_build_device kwarg threaded through. Checkpoint forward-compat note: new checkpoints without scaler state cannot be loaded by old code.
examples/cfd/external_aerodynamics/globe/drivaer/train.py	Identical GradScaler removal and tree_build_device addition as airfrans/train.py; same checkpoint forward-compat concern applies.
test/models/globe/test_barnes_hut_kernel.py	Removes tests for now-deleted leaf_node_ids/leaf_seg_ids attributes; adds precision regression guards for the fp64 cumsum fix, a bf16 autocast forward correctness test, and a CUDA sync-budget regression test. Coverage is well-targeted at the riskiest new code paths.
pyproject.toml	Bumps tensordict floor to 0.12.2 and adds psutil>=6.0.0 for the CPU memory-query fallback path. Both additions are justified by the PR changes.

Comments Outside Diff (1)

1. examples/cfd/external_aerodynamics/globe/airfrans/train.py, line 591-597 (link)

   Checkpoint forward-compatibility break

   Checkpoints saved with the new code (no scaler key) will fail to load with the old training script, which still calls load_checkpoint(..., scaler=scaler, ...) and expects that key in the saved state dict. While the PR guarantees backward compat (old checkpoint → new code), users who need to roll back (e.g., a bugfix release reverting to the old code) or who share a checkpoint with a colleague running an older version will hit a hard failure. Worth documenting this one-way migration in the training README or a comment near save_checkpoint.

_{Reviews (1): Last reviewed commit: "Squashed commit of the following:" | Re-trigger Greptile}