Skip long-lived objects in the per-step gc.collect() of QAT and QAD

CHANGELOG.rst

@@ -23,6 +23,7 @@ Changelog
 - Add a fake-quantized weight cache for QAT and QAD. Enroll the model with ``mtq.utils.enable_weight_cache(model)`` and wrap the forward and backward passes of each optimizer step in ``mtq.utils.weight_cache_scope(model)``: each weight is then fake-quantized once per step instead of once per micro-batch and activation recomputation, at the cost of one extra copy of each cached weight. It needs a straight-through weight gradient (``pass_through_bwd=True``, the default).
 - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2.
 - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration.
+- Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B.
 
 *Speculative Decoding*
 

examples/llm_qat/ARGUMENTS.md

@@ -65,6 +65,6 @@ Extends [HuggingFace TrainingArguments](https://huggingface.co/docs/transformers
 | `--trainable_params` | `list[str]` | `None` | Glob patterns (fnmatch) for parameters that should be trainable. All other parameters will be frozen. Mutually exclusive with frozen_params. |
 | `--frozen_params` | `list[str]` | `None` | Glob patterns (fnmatch) for parameters that should be frozen. Mutually exclusive with trainable_params. |
 | `--lr_config` | `str` | `None` | Path to a YAML file mapping fnmatch patterns to optimizer kwargs (e.g. lr, weight_decay). First matching pattern wins per parameter. See examples/llm_qat/configs/train/lr/lr_config_example.yaml. |
-| `--manual_gc` | `bool` | `False` | Run `gc.collect()` before each training/prediction step to work around GPU memory leaks during QAT/distillation. |
+| `--manual_gc` | `bool` | `False` | Run `gc.collect()` before each training/prediction step to work around GPU memory leaks during QAT/distillation. After the first optimizer step, objects that live until `train()` returns are frozen (`gc.freeze()`) so these collections skip them. |
 | `--liger_ce_label_smoothing` | `float` | `0.0` | Label smoothing for Liger fused CE loss. Only used when --use_liger_kernel is enabled. |
 | `--lora` | `bool` | `False` | Whether to add LoRA (Low-Rank Adaptation) adapter before training. When using real quantization, the LoRA adapter must be set, as quantized weights will be frozen during training. |

modelopt/torch/opt/plugins/transformers.py

@@ -256,7 +256,9 @@ class ModelOptTrainerArguments(ModelOptHFArguments):
         metadata={
             "help": (
                 "Run `gc.collect()` before each training/prediction step to work around "
-                "GPU memory leaks during QAT/distillation."
+                "GPU memory leaks during QAT/distillation. After the first optimizer step, objects "
+                "that live until `train()` returns are frozen (`gc.freeze()`) so these "
+                "collections skip them."
             ),
         },
     )
@@ -526,6 +528,7 @@ def __init__(
         if trainer_args is None and isinstance(self.args, ModelOptTrainerArguments):
             trainer_args = self.args
         self.trainer_args = trainer_args or ModelOptTrainerArguments()
+        self._gc_frozen = False
         self._lr_config = self._resolve_lr_config(lr_config, self.trainer_args)
         self._apply_gradient_checkpointing_defaults()
         self.add_callback(_MemoryReportCallback())
@@ -554,15 +557,38 @@ def _prepare_model(self, model):
     def training_step(self, *args, **kwargs):
         """Run gc.collect() before the training step if manual_gc is enabled."""
         if self.trainer_args.manual_gc:
-            gc.collect()
+            self._collect_garbage(freeze=True)
         return super().training_step(*args, **kwargs)
 
     def prediction_step(self, *args, **kwargs):
         """Run gc.collect() before the prediction step if manual_gc is enabled."""
         if self.trainer_args.manual_gc:
-            gc.collect()
+            self._collect_garbage()
         return super().prediction_step(*args, **kwargs)
 
+    def _collect_garbage(self, freeze=False):
+        """``gc.collect()``; with ``freeze``, also ``gc.freeze()`` once training is under way.
+
+        After the first optimizer step the models, quantizers and optimizer state exist and stay
+        alive until ``train()`` returns. Frozen, later collections skip them instead of rescanning
+        them every step, which took about 0.5 s per collection on a 30B MoE.
+        """
+        gc.collect()
+        if freeze and not self._gc_frozen and self.state.global_step >= 1:
+            gc.freeze()
+            self._gc_frozen = True
+
+    def train(self, *args, **kwargs):
+        """Train, then unfreeze what the per-step collections froze so it can be freed."""
+        try:
+            return super().train(*args, **kwargs)
+        finally:
+            if self._gc_frozen:
+                # Also unfreezes anything frozen before (Python itself starts with a few frozen
+                # objects); that only makes later collections scan them again.
+                gc.unfreeze()
+                self._gc_frozen = False
+
     def _load_best_model(self, *args, **kwargs):
         """Run gc.collect() before loading the best model if manual_gc is enabled."""
         if self.trainer_args.manual_gc:

modelopt/torch/quantization/plugins/transformers_trainer.py

@@ -256,14 +256,14 @@ def forward_loop(model):
 
     def training_step(self, *args, **kwargs):
         """Training step."""
-        gc.collect()
+        self._collect_garbage(freeze=True)
         if self.quant_cfg is not None and not is_quantized(self.model):
             self._quantize_model()
         return super().training_step(*args, **kwargs)
 
     def prediction_step(self, *args, **kwargs):
         """Prediction step."""
-        gc.collect()
+        self._collect_garbage()
         if self.quant_cfg is not None and not is_quantized(self.model):
             self._quantize_model()
         return super().prediction_step(*args, **kwargs)

tests/unit/torch/opt/plugins/test_trainer_gc.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 gc
+import weakref
+
+import pytest
+import torch
+from torch import nn
+
+transformers = pytest.importorskip("transformers")
+
+from modelopt.torch.opt.plugins.transformers import ModelOptHFTrainer, ModelOptTrainerArguments
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Embedding(32, 8)
+        self.lm_head = nn.Linear(8, 32)
+
+    def forward(self, input_ids, labels=None, **kwargs):
+        logits = self.lm_head(self.embed(input_ids))
+        loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
+        return {"loss": loss, "logits": logits}
+
+
+class _FreezeCounts(transformers.TrainerCallback):
+    def __init__(self):
+        self.after_step = {}
+
+    def on_step_end(self, args, state, control, **kwargs):
+        self.after_step[state.global_step] = gc.get_freeze_count()
+
+
+def test_manual_gc_freezes_during_train_and_unfreezes_after(tmp_path):
+    """``manual_gc`` freezes what lives through training and ``train()`` unfreezes it.
+
+    Collections skip frozen objects, which is the speed-up. A frozen object in a reference cycle is
+    never collected, so without the unfreeze the model would outlive the trainer.
+    """
+    frozen_before = gc.get_freeze_count()  # Python itself starts with a few hundred frozen objects
+    freeze_counts = _FreezeCounts()
+    model = _TinyModel()
+    model.__dict__["_cycle"] = [model]  # only the cycle collector can free it
+    model_ref = weakref.ref(model)
+    dataset = [
+        {"input_ids": torch.randint(0, 32, (8,)), "labels": torch.randint(0, 32, (8,))}
+        for _ in range(8)
+    ]
+    args = transformers.TrainingArguments(
+        output_dir=str(tmp_path),
+        max_steps=3,
+        per_device_train_batch_size=2,
+        report_to="none",
+        use_cpu=True,
+    )
+    trainer = ModelOptHFTrainer(
+        model=model,
+        args=args,
+        trainer_args=ModelOptTrainerArguments(manual_gc=True),
+        train_dataset=dataset,
+        callbacks=[freeze_counts],
+    )
+    try:
+        trainer.train()
+        # The collection before the second step freezes the model and optimizer state.
+        assert all(freeze_counts.after_step[step] > frozen_before for step in (2, 3))
+        # gc.unfreeze() also returns the objects frozen before training.
+        assert gc.get_freeze_count() == 0
+        del trainer, model
+        gc.collect()
+        assert model_ref() is None
+    finally:
+        gc.unfreeze()