From d941a3befa1fe06ba9e33013c603c5ffd703e682 Mon Sep 17 00:00:00 2001
From: alisonshao <54658187+alisonshao@users.noreply.github.com>
Date: Wed, 26 Nov 2025 18:27:51 -0800
Subject: [PATCH] Fix nightly test failure: NSA indexer dtype (#14017)

---
 test/nightly/test_nsa_indexer.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/test/nightly/test_nsa_indexer.py b/test/nightly/test_nsa_indexer.py
index e6c47c909..a0e953625 100644
--- a/test/nightly/test_nsa_indexer.py
+++ b/test/nightly/test_nsa_indexer.py
@@ -248,12 +248,15 @@ class TestNSAIndexer(CustomTestCase):
         # Move indexer to CUDA device
         indexer = indexer.to(device=self.device)
 
-        # Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32)
+        # Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32
+        # and weights_proj's float32 - it uses params_dtype=torch.float32 in production)
         # Need to recursively convert LinearBase submodules (like ReplicatedLinear)
         for name, module in indexer.named_modules():
             # Check for LinearBase (parent of ReplicatedLinear) but exclude LayerNorm
+            # Also exclude weights_proj which uses float32 params in production
             if isinstance(module, LinearBase) and not isinstance(module, LayerNorm):
-                module.to(dtype=self.dtype)
+                if "weights_proj" not in name:
+                    module.to(dtype=self.dtype)
 
         return indexer