From d92fac231e94b51338bafcb5ae9987f6d2ce0511 Mon Sep 17 00:00:00 2001
From: Akihiro Nitta <nitta@akihironitta.com>
Date: Tue, 3 Nov 2026 07:26:06 +0000
Subject: [PATCH] Reduce TabFM row embedding peak memory

Build the column projection and row norms of the TabFM row embedding
from `sdm.nn.RMSNorm`, so that under autocast they write the autocast
dtype instead of float32 copies of the whole embedding.
---
 sdm/models/tabfm/row_embedding.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/sdm/models/tabfm/row_embedding.py b/sdm/models/tabfm/row_embedding.py
index d49a7d22d4..3ae887f491 100644
--- a/sdm/models/tabfm/row_embedding.py
+++ b/sdm/models/tabfm/row_embedding.py
@@ -27,14 +27,13 @@
     Linear,
     ModuleList,
     Parameter,
-    RMSNorm,
     Sequential,
 )
 
 from sdm.cache import Cache, KVCacheEntry
 from sdm.models.tabfm.block import TabFMTransformerBlock
 from sdm.models.tabfm.cell_embedding import CellEmbedding
-from sdm.nn import InducedTransformerBlock, RotaryEmbedding
+from sdm.nn import InducedTransformerBlock, RMSNorm, RotaryEmbedding
 
 
 class RowEmbedding(torch.nn.Module):