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..e5502569b61 --- /dev/null +++ b/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()