From 466ff20e51489883625a9b11e832fe7775d2c88e Mon Sep 17 00:00:00 2001 From: Serge Panev Date: Tue, 17 Mar 2026 09:47:58 -0700 Subject: [PATCH] [Model] Fix NemotronH OOM on unified-mem systems: stream weights + safetensors cleanup (#20580) Signed-off-by: Serge Panev --- python/sglang/srt/models/nemotron_h.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 669eb4fe5..2a83158c3 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -774,12 +774,6 @@ class NemotronHForCausalLM(nn.Module): def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False ) -> None: - updated_weights = [] - for name, loaded_weight in weights: - name = replace_prefix(name, self.remap_prefix) - name = replace_substrings(name, self.remap_substr) - updated_weights.append((name, loaded_weight)) - # - FusedMoe.w1 (aka gate_proj) should be up_proj since that's # what the activation is applied to # - FusedMoe.w3 (aka up_proj) should be ignored since we're @@ -793,7 +787,13 @@ class NemotronHForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) - for name, loaded_weight in updated_weights: + # Stream weights directly from the generator to avoid buffering + # the entire checkpoint (~75 GB) into a Python list. On unified- + # memory systems (e.g. DGX Spark, 119 GB) the old buffered path + # caused OOM: skeleton 81.6 GB + buffer 75 GB = 157 GB peak. + for name, loaded_weight in weights: + name = replace_prefix(name, self.remap_prefix) + name = replace_substrings(name, self.remap_substr) if is_mtp: if "mtp" not in name: continue