From 52ebb8d1e42d00ed099d70a32332aeae37c0a191 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Wed, 28 Oct 2026 18:28:04 -0700 Subject: [PATCH 1/2] Skip long-lived objects in the per-step gc.collect() of QAT and QAD QATTrainer runs gc.collect() before every training step, and manual_gc runs another one. Each walks every object the cycle collector tracks: about 1.9M on Qwen3-30B-A3B QAD (models, quantizers, optimizer state), 0.5 s per collection and 1.1 s of a 3.6 s step. From the second training step on, the trainer freezes what is alive after the collection once (gc.freeze()), so later collections skip it, and unfreezes it when train() returns so the model can be freed. Python 3.12 starts with a few hundred frozen objects, so the trainer tracks its own freeze instead of checking gc.get_freeze_count(). Co-Authored-By: Claude Opus 5.5 Signed-off-by: Shiyang Chen --- CHANGELOG.rst | 1 + examples/llm_qat/ARGUMENTS.md | 2 +- modelopt/torch/opt/plugins/transformers.py | 32 ++++++++- .../plugins/transformers_trainer.py | 4 +- .../unit/torch/opt/plugins/test_trainer_gc.py | 71 +++++++++++++++++++ 5 files changed, 104 insertions(+), 6 deletions(-) create mode 100644 tests/unit/torch/opt/plugins/test_trainer_gc.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9a267cba281..37d9036368a 100755 --- a/CHANGELOG.rst +++ b/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* diff --git a/examples/llm_qat/ARGUMENTS.md b/examples/llm_qat/ARGUMENTS.md index 8e1022d49dc..b59984e1e85 100644 --- a/examples/llm_qat/ARGUMENTS.md +++ b/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. | diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index 3ebef890f56..83155be5344 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/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: diff --git a/modelopt/torch/quantization/plugins/transformers_trainer.py b/modelopt/torch/quantization/plugins/transformers_trainer.py index b3d5619beb2..764fa472d13 100644 --- a/modelopt/torch/quantization/plugins/transformers_trainer.py +++ b/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) diff --git a/tests/unit/torch/opt/plugins/test_trainer_gc.py b/tests/unit/torch/opt/plugins/test_trainer_gc.py new file mode 100644 index 00000000000..7a0d64143df --- /dev/null +++ b/tests/unit/torch/opt/plugins/test_trainer_gc.py @@ -0,0 +1,71 @@ +# 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} + + +def test_model_is_freed_after_train_with_manual_gc(tmp_path): + """``manual_gc`` freezes what lives through training; ``train()`` must unfreeze it. + + A frozen object in a reference cycle is never collected, so the model would outlive the trainer. + """ + 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, + ) + try: + trainer.train() + del trainer, model + gc.collect() + assert model_ref() is None + finally: + gc.unfreeze() From 5539b116e134c106e0accb8818a6e0e39b258fd5 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Thu, 29 Oct 2026 17:08:21 -0700 Subject: [PATCH 2/2] Test that manual_gc freezes during train(), not only that it unfreezes The test passed with gc.freeze() removed, so nothing guarded the speed-up. A callback now records gc.get_freeze_count() after each step: it must be above its value before train() from the second step on, and zero once train() returns, since gc.unfreeze() also returns what Python froze at startup. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Shiyang Chen --- .../unit/torch/opt/plugins/test_trainer_gc.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/opt/plugins/test_trainer_gc.py b/tests/unit/torch/opt/plugins/test_trainer_gc.py index 7a0d64143df..e5502569b61 100644 --- a/tests/unit/torch/opt/plugins/test_trainer_gc.py +++ b/tests/unit/torch/opt/plugins/test_trainer_gc.py @@ -37,11 +37,22 @@ def forward(self, input_ids, labels=None, **kwargs): return {"loss": loss, "logits": logits} -def test_model_is_freed_after_train_with_manual_gc(tmp_path): - """``manual_gc`` freezes what lives through training; ``train()`` must unfreeze it. +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. - A frozen object in a reference cycle is never collected, so the model would outlive the trainer. + 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) @@ -61,9 +72,14 @@ def test_model_is_freed_after_train_with_manual_gc(tmp_path): 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