[akihironitta]

Summary

Build the TabFM row embedding's column-projection and row norms from sdm.nn.RMSNorm instead of torch.nn.RMSNorm. This cuts TabFM's peak GPU memory at 50k context rows by about 40%, with identical predictions.

Description of the change

Under CUDA autocast, torch.nn.RMSNorm runs in float32. So on the inference path, x.copy_(col_proj(x)) and x.copy_(row_norm(x)) in RowEmbedding create float32 copies of the whole embedding before writing it back. sdm.nn.RMSNorm, which the TabFM transformer blocks already use, writes the autocast dtype directly. It subclasses torch.nn.RMSNorm, so parameters and checkpoint keys are unchanged.

Validation

H100 PCIe 80GB, bf16 TabFM under bf16 autocast, 32 estimators, synthetic inputs. Main and this branch ran back to back on the same GPU. Peak is torch.cuda.max_memory_allocated(), including weights.

Context / query rows, features	Task	Peak, main (MiB)	Peak, this PR (MiB)
50,000 / 10,000, 100	classification	22,989.5	13,472.6
50,000 / 10,000, 100	regression	22,997.5	13,480.6

Predictions are identical to main in all six cases measured: 1k, 10k and 50k context rows, for both classification and regression. At 1k and 10k rows, peak memory and latency are unchanged.

[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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.33%. Comparing base (c51bc37) to head (d92fac2).
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1331   +/-   ##
=======================================
  Coverage   82.33%   82.33%           
=======================================
  Files         185      185           
  Lines       11794    11794           
=======================================
  Hits         9710     9710           
  Misses       2084     2084           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.