Add step-aligned garbage collection to SALMAutomodel

Author: pzelasko

docs/source/speechlm2/configs.rst

@@ -170,6 +170,22 @@ SALMAutomodel-Specific Options
 The SALMAutomodel config exposes a few extra knobs that pass through to NeMo
 Automodel. All are optional — defaults preserve standard behavior.
 
+**Garbage collection:**
+
+.. code-block:: yaml
+
+    model:
+      # Optional positive optimizer-step interval; null keeps automatic GC.
+      gc_every_steps: null
+
+Setting ``gc_every_steps`` to a positive integer disables Python's automatic
+garbage collector at fit start and uses NeMo Automodel's generation-1 collector
+at that optimizer-step cadence. This avoids an occasional generation-2 scan on
+one distributed rank delaying all peers at the next collective. The cadence is
+counted in optimizer steps, so gradient accumulation does not increase the
+collection frequency. Leave it ``null`` unless profiling shows GC-related rank
+stragglers.
+
 **MoE training:**
 
 .. code-block:: yaml

examples/speechlm2/conf/salm_automodel.yaml

@@ -14,6 +14,11 @@ model:
   # Set to true to use SALMAutomodel (NeMo Automodel backend) instead of SALM (HF Transformers backend).
   use_nemo_automodel: true
 
+  # Optional: disable asynchronous full-heap Python GC and run deterministic
+  # generation-1 collections every N optimizer steps. null preserves Python's
+  # default automatic GC behavior.
+  gc_every_steps: null
+
   # Regexp (re.compile) patterns matching parameters to be frozen.
   freeze_params:
     # Frozen LLM (embed_tokens stays inside llm, so this pattern covers it too)

examples/speechlm2/conf/salm_automodel_pee.yaml

@@ -14,6 +14,11 @@ model:
   # Set to true to use SALMAutomodel (NeMo Automodel backend) instead of SALM (HF Transformers backend).
   use_nemo_automodel: true
 
+  # Optional: disable asynchronous full-heap Python GC and run deterministic
+  # generation-1 collections every N optimizer steps. null preserves Python's
+  # default automatic GC behavior.
+  gc_every_steps: null
+
   # Regexp (re.compile) patterns matching parameters to be frozen.
   # PEE recipe: freeze the LLM and the Sortformer diarizer expert; keep the ASR
   # Conformer encoder (perception.encoder.asr_encoder) and the fusion layers

nemo/collections/speechlm2/models/salm_automodel.py

@@ -32,6 +32,7 @@
 from nemo.collections.speechlm2.models.salm import _resolve_audios_in_prompt, replace_placeholders_and_build_targets
 from nemo.collections.speechlm2.parts.automodel_lora import ensure_lora_trainable, make_peft_config, maybe_install_lora
 from nemo.collections.speechlm2.parts.encoder_chunking import encode_audio_with_optional_chunking
+from nemo.collections.speechlm2.parts.gc import GarbageCollectionManager
 from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin
 from nemo.collections.speechlm2.parts.multispeaker import build_speaker_tokens, maybe_init_lss_loss
 from nemo.collections.speechlm2.parts.optim_setup import configure_optimizers, is_frozen
@@ -69,6 +70,7 @@ def __init__(self, cfg) -> None:
 
         self._use_fsdp = False
         self._use_tp = False
+        self._garbage_collection = GarbageCollectionManager(self.cfg.get("gc_every_steps", None))
 
         if self.cfg.get("init_configure_model", False):
             self.configure_model()
@@ -379,6 +381,12 @@ def on_fit_start(self) -> None:
         averaging (see ``_configure_moe_aux_loss_scaler``)."""
         self._validate_parallelism_compatibility()
         self._configure_moe_aux_loss_scaler()
+        self._garbage_collection.on_fit_start()
+
+    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure=None) -> None:
+        """Run configured manual GC after each completed optimizer step."""
+        super().optimizer_step(epoch, batch_idx, optimizer, optimizer_closure)
+        self._garbage_collection.on_optimizer_step()
 
     def _validate_parallelism_compatibility(self) -> None:
         """Raise on known-incompatible THD/CP/backend configurations.

nemo/collections/speechlm2/parts/gc.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.utils import logging
+
+
+class GarbageCollectionManager:
+    """Manage deterministic Python garbage collection during distributed training.
+
+    When enabled, automatic garbage collection is replaced at fit start by
+    NeMo Automodel's generation-1 collector. The manager owns the optimizer-step
+    counter so model implementations only need to forward lifecycle events.
+    """
+
+    def __init__(self, gc_every_steps: int | None) -> None:
+        if gc_every_steps is not None and (
+            isinstance(gc_every_steps, bool)
+            or not isinstance(gc_every_steps, int)
+            or gc_every_steps <= 0
+        ):
+            raise ValueError(
+                f"model.gc_every_steps must be a positive integer or null, got {gc_every_steps!r}"
+            )
+        self.gc_every_steps = gc_every_steps
+        self._collector = None
+        self._optimizer_step_count = 0
+
+    def on_fit_start(self) -> None:
+        """Disable automatic GC and initialize the configured manual collector."""
+        if self.gc_every_steps is None:
+            return
+
+        from nemo_automodel.components.training.garbage_collection import (
+            GarbageCollection,
+        )
+
+        self._collector = GarbageCollection(gc_every_steps=self.gc_every_steps)
+        self._optimizer_step_count = 0
+        logging.info(
+            "Automatic Python GC disabled; generation-1 collection will run every %d optimizer steps",
+            self.gc_every_steps,
+        )
+
+    def on_optimizer_step(self) -> None:
+        """Advance the manual collector after a completed optimizer step."""
+        if self._collector is None:
+            return
+        self._optimizer_step_count += 1
+        self._collector.run(self._optimizer_step_count)

tests/collections/speechlm2/test_gc.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from nemo.collections.speechlm2.parts.gc import GarbageCollectionManager
+
+
+def test_garbage_collection_manager_owns_step_state(monkeypatch):
+    calls = []
+
+    class FakeGarbageCollection:
+        def __init__(self, gc_every_steps):
+            calls.append(("init", gc_every_steps))
+
+        def run(self, step_count):
+            calls.append(("run", step_count))
+
+    import nemo_automodel.components.training.garbage_collection as gc_module
+
+    monkeypatch.setattr(gc_module, "GarbageCollection", FakeGarbageCollection)
+    manager = GarbageCollectionManager(gc_every_steps=10)
+
+    manager.on_fit_start()
+    manager.on_optimizer_step()
+    manager.on_optimizer_step()
+
+    assert calls == [("init", 10), ("run", 1), ("run", 2)]
+
+
+def test_garbage_collection_manager_is_noop_when_disabled():
+    manager = GarbageCollectionManager(gc_every_steps=None)
+
+    manager.on_fit_start()
+    manager.on_optimizer_step()
+
+    assert manager._collector is None
+    assert manager._optimizer_step_count == 0
+
+
+@pytest.mark.parametrize("value", [True, False, 0, -1, 1.5, "10"])
+def test_garbage_collection_manager_rejects_invalid_interval(value):
+    with pytest.raises(ValueError, match="gc_every_steps"):
+        GarbageCollectionManager(gc_every_steps=value)

tests/collections/speechlm2/test_salm_automodel.py

@@ -18,6 +18,7 @@
 import torch
 from lhotse import CutSet, SupervisionSegment
 from lhotse.testing.dummies import dummy_cut, dummy_recording
+from lightning import LightningModule
 from transformers import GenerationConfig
 
 from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
@@ -205,6 +206,25 @@ def test_salm_automodel_training_step_uses_dataloader_iter_signature():
     assert list(inspect.signature(SALMAutomodel.training_step).parameters) == ["self", "dataloader_iter"]
 
 
+def test_salm_automodel_notifies_garbage_collection_after_optimizer_step(monkeypatch):
+    calls = []
+
+    class FakeGarbageCollectionManager:
+        def on_optimizer_step(self):
+            calls.append("gc")
+
+    model = SALMAutomodel.__new__(SALMAutomodel)
+    torch.nn.Module.__init__(model)
+    model._garbage_collection = FakeGarbageCollectionManager()
+    monkeypatch.setattr(
+        LightningModule,
+        "optimizer_step",
+        lambda *args, **kwargs: calls.append("optimizer"),
+    )
+    model.optimizer_step(0, 0, object())
+    assert calls == ["optimizer", "gc"]
+
+
 def test_salm_automodel_record_training_stats_uses_thd_metadata():
     model = SALMAutomodel.__new__(SALMAutomodel)
     batch = {"input_ids": torch.zeros(3, 7, dtype=torch.long)}