trending/Model-Optimizer
1
0
Fork 0
mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-02 03:14:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge puzzletron compression algorithm (#1121)

Browse Source 
### What does this PR do?

Implement puzzletron compression algorithm based on Puzzle paper
(https://arxiv.org/abs/2411.19146)

<details>
<summary> Th list of reviewed and merged MRs that resulted in the
feature/puzzletron branch</summary>

Merging dkorzekwa/any_model to feature/puzzletron

[Add anymodel directories to feature/puzzletron by danielkorzekwa · Pull
Request #974 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/974)
- merged

[Draft: anymodel activation scoring by danielkorzekwa · Pull Request
#989 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/989)
- merged

[Draft: Merge anymodel pruning by danielkorzekwa · Pull Request #990 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/990/)
- merged

[Draft: Merging anymodel:build_library_and_stats by danielkorzekwa ·
Pull Request #993 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/993)
- merged

[Dkorzekwa/any model calc one block scores by danielkorzekwa · Pull
Request #994 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/994)
- merged

[Draft: merge any_model: mip_and_realize_models by danielkorzekwa · Pull
Request #995 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/995)
- merged

[Dkorzekwa/any model other modeqls by danielkorztiekwa · Pull Request
#1007 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1007/)
- merged

PR to 1007: https://github.com/NVIDIA/Model-Optimizer/pull/1039 - merged

[Dkorzekwa/anymodel gptoss by danielkorzekwa · Pull Request #1020 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1020)
- merged

[Merge any_model tutorial by danielkorzekwa · Pull Request #1035 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1035)
- merged

[Merge mbridge distillation for any_model by danielkorzekwa · Pull
Request #1036 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1036)
- merged

[MR branch for the remaining difference between dkorzekwa/any_model an…
by danielkorzekwa · Pull Request #1047 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1047)
- merged

[Dkorzekwa/decilm hf code cleanup by danielkorzekwa · Pull Request #1071
·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1071)
- merged

[Dkorzekwa/decilm hf code cleanup 2 by danielkorzekwa · Pull Request
#1073 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1073)
- merged

[Dkorzekwa/anymodel subblock stats by danielkorzekwa · Pull Request
#1085 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1085)
- merged

[Dkorzekwa/anymodel subblock stats nodecilm by danielkorzekwa · Pull
Request #1102 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1102)
- merged

[Dkorzekwa/decilm cleanup post subblockstats by danielkorzekwa · Pull
Request #1103 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1103)
- merged

[code clean up by danielkorzekwa · Pull Request #1110 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1110)
- merged

Merging into main:

[Activation hooks redesign (reuse hooks component across both minitron
and puzzletron) by danielkorzekwa · Pull Request #1022 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1022)
- merged

[Dkorzekwa/puzzletron use importance hooks from prune by danielkorzekwa
· Pull Request #1115 ·
NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1115)
- merged

</details>

<!-- Details about the change. -->

### Usage

Puzzletron tutorial:

https://github.com/NVIDIA/Model-Optimizer/tree/feature/puzzletron/examples/puzzletron

### Testing
The main e2e test for compressing 9 models with Puzzletron:

https://github.com/NVIDIA/Model-Optimizer/blob/feature/puzzletron/tests/gpu/torch/puzzletron/test_puzzletron.py

2-gpu nightly tests: 

-
https://github.com/NVIDIA/Model-Optimizer/actions/runs/24468209205/job/71501061203
-
https://github.com/NVIDIA/Model-Optimizer/actions/runs/24470214159/job/71508152952

### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added Puzzletron: end-to-end heterogeneous pruning & NAS workflow with
AnyModel support, example pipelines, deployment and evaluation
utilities, and tools for converting/pruning and exporting compressed
checkpoints.

* **Documentation**
* Comprehensive Puzzletron tutorials, model-specific guides, evaluator
instructions, example configs, and changelog entry.

* **Chores**
* CI/workflow updates (extras installation, longer GPU test timeout),
pre-commit hook exclusion updated, and CODEOWNERS entries added.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Liana Mikaelyan <lmikaelyan@nvidia.com>
Signed-off-by: Liana Mikaelyan <45925959+LianaMikael@users.noreply.github.com>
Signed-off-by: Daniel Korzekwa <daniel.korzekwa@gmail.com>
Signed-off-by: jrausch <jrausch@nvidia.com>
Signed-off-by: root <root@pool0-00848.cm.cluster>
Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Liana Mikaelyan <lmikaelyan@nvidia.com>
Co-authored-by: Liana Mikaelyan <45925959+LianaMikael@users.noreply.github.com>
Co-authored-by: J Rausch <38429553+j-rausch@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Daniel Korzekwa
2026-04-16 00:48:11 +05:30
committed by GitHub
co-authored by Keval Morabia Liana Mikaelyan Liana Mikaelyan J Rausch Claude Sonnet 4.6
parent dec2952992
commit 361f7e391b
235 changed files with 24885 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/CODEOWNERS
+2
View File
@@ -24,6 +24,7 @@ modelopt/torch/nas @NVIDIA/modelopt-torch-nas-prune-codeowners
modelopt/torch/opt @NVIDIA/modelopt-torch-opt-codeowners
modelopt/torch/peft @NVIDIA/modelopt-torch-peft-codeowners
modelopt/torch/prune @NVIDIA/modelopt-torch-nas-prune-codeowners
modelopt/torch/puzzletron @NVIDIA/modelopt-torch-puzzletron-codeowners
modelopt/torch/quantization @NVIDIA/modelopt-torch-quantization-codeowners
modelopt/torch/sparsity @NVIDIA/modelopt-torch-sparsity-codeowners
modelopt/torch/speculative @NVIDIA/modelopt-torch-speculative-codeowners
@@ -49,6 +50,7 @@ modelopt_recipes @NVIDIA/modelopt-recipes-codeowners
/examples/model_hub @NVIDIA/modelopt-examples-model_hub-codeowners
/examples/onnx_ptq @NVIDIA/modelopt-onnx-codeowners
/examples/pruning @NVIDIA/modelopt-torch-nas-prune-codeowners
/examples/puzzletron @NVIDIA/modelopt-torch-puzzletron-codeowners
/examples/specdec_bench @NVIDIA/modelopt-torch-speculative-codeowners
/examples/speculative_decoding @NVIDIA/modelopt-torch-speculative-codeowners
/examples/torch_onnx @NVIDIA/modelopt-onnx-codeowners
.github/workflows/_example_tests_runner.yml
+2 -1
View File
@@ -48,6 +48,7 @@ jobs:
- name: Install dependencies
run: |
# use `python -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers
pip uninstall -y nvidia-modelopt
python -m pip install ".${{ inputs.pip_install_extras }}"
if [[ "${{ inputs.example }}" == *"diffusers"* ]]; then
@@ -64,7 +65,7 @@ jobs:
COVERAGE_FILE: ${{ github.workspace }}/.coverage
run: |
echo "Running tests for: ${{ inputs.example }}"
pytest tests/examples/${{ inputs.example }} --cov
python -m pytest tests/examples/${{ inputs.example }} --cov
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v5
with:
.github/workflows/example_tests.yml
+2 -2
View File
@@ -132,7 +132,7 @@ jobs:
docker_image: "nvcr.io/nvidia/nemo:26.02"
example: ${{ matrix.example }}
timeout_minutes: 30
pip_install_extras: "[hf,dev-test]"
pip_install_extras: "[hf,puzzletron,dev-test]"
runner: linux-amd64-gpu-rtxpro6000-latest-1
nemo-non-pr:
@@ -144,7 +144,7 @@ jobs:
docker_image: "nvcr.io/nvidia/nemo:26.02"
example: ${{ matrix.example }}
timeout_minutes: 30
pip_install_extras: "[hf,dev-test]"
pip_install_extras: "[hf,puzzletron,dev-test]"
runner: linux-amd64-gpu-rtxpro6000-latest-2
##### ONNX/TensorRT Example Tests #####
.github/workflows/gpu_tests.yml
+1 -1
View File
@@ -68,7 +68,7 @@ jobs:
matrix:
include:
- example: gpu
timeout: 45
timeout: 60
container_image: pytorch:26.01-py3
# tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers until https://github.com/tbenthompson/cppimport/pull/98
- example: gpu-regression
.pre-commit-config.yaml
+1
View File
@@ -94,6 +94,7 @@ repos:
modelopt/onnx/quantization/ort_patching.py|
modelopt/torch/_deploy/utils/onnx_utils.py|
modelopt/torch/export/transformer_engine.py|
modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py|
modelopt/torch/quantization/export_onnx.py|
modelopt/torch/quantization/plugins/attention.py|
modelopt/torch/sparsity/attention_sparsity/methods/vsa_utils.py|
CHANGELOG.rst
+1
View File
@@ -7,6 +7,7 @@ Changelog
**New Features**
- Support full Transformer Engine spec for Minitron pruning (``mcore_minitron``). Now we no longer need to use custom ModelOpt spec. Note that this does not affect the usage of the pruning workflow but makes pruning slightly faster and may result in slightly different pruned model because of different kernel and numerics.
- Add Puzzletron - a new algorithm for heterogeneous pruning of LLM and VLM models. See `examples/puzzletron/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/puzzletron>`_ for more details.
- Added iterator interface using CalibrationDataReader in ONNX quantization workflow.
- Add N:M sparse softmax support to the Triton flash attention kernel (``modelopt.torch.kernels.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_ for usage.
- Add skip-softmax skipping to the Triton flash attention kernel (``modelopt.torch.kernels.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_ for usage.
docs/source/conf.py
+9
View File
@@ -31,6 +31,7 @@
# import sys
# sys.path.insert(0, os.path.abspath('.'))
import contextlib
import os
import sys
@@ -44,6 +45,14 @@ from modelopt import __version__
sys.path.insert(0, os.path.abspath("../../"))
sys.path.append(os.path.abspath("./_ext"))
# Pre-import modelopt.torch so it is cached in sys.modules before Sphinx applies
# autodoc_mock_imports. Mocking triton/tensorrt_llm at the Sphinx level can break
# transitive imports (transformers, transformer_engine, …) and cause modelopt.torch
# to fail inside autosummary. Importing here — while the real packages are still on
# sys.path — avoids that problem entirely.
with contextlib.suppress(Exception):
import modelopt.torch # noqa: F401
# -- Project information -----------------------------------------------------
project = "Model Optimizer" # pylint: disable=C0103
examples/llm_eval/README.md
+16
View File
@@ -40,6 +40,22 @@ accelerate launch --multi_gpu --num_processes <num_copies_of_your_model> \
--batch_size 4
```
### Heterogeneous Pruned Checkpoints (Puzzletron)
Heterogeneous pruned checkpoints produced by Puzzletron are automatically detected and loaded with the appropriate model patcher. No additional flags are needed beyond specifying the checkpoint path:
```sh
python lm_eval_hf.py --model hf \
--model_args pretrained=path/to/anymodel/checkpoint,dtype=bfloat16,parallelize=True \
--tasks mmlu \
--num_fewshot 5 \
--batch_size 4
```
For a quick smoke test, add `--limit 10`.
> **Note:** Requires the `puzzletron` extra to be installed (`pip install -e ".[puzzletron]"`).
### Quantized (simulated)
- For simulated quantization with any of the default quantization formats:
examples/llm_eval/lm_eval_hf.py
+51 -2
View File
@@ -36,11 +36,19 @@
# 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 contextlib
import warnings
import datasets
import lm_eval
from lm_eval import utils
from lm_eval.__main__ import cli_evaluate, parse_eval_args, setup_parser
if not lm_eval.__version__.startswith("0.4.8"):
warnings.warn(
f"lm_eval_hf.py is tested with lm-eval 0.4.8; found {lm_eval.__version__}. "
"Later versions may have incompatible API changes."
)
from lm_eval.api.model import T
from lm_eval.models.huggingface import HFLM
from quantization_utils import quantize_model
@@ -50,9 +58,29 @@ import modelopt.torch.opt as mto
from modelopt.torch.quantization.utils import is_quantized
from modelopt.torch.sparsity.attention_sparsity.conversion import is_attn_sparsified
try:
import modelopt.torch.puzzletron as mtpz
_ANYMODEL_AVAILABLE = True
except ImportError:
_ANYMODEL_AVAILABLE = False
def _anymodel_patcher_context(pretrained, trust_remote_code=False):
"""Return a deci_x_patcher context if *pretrained* is a Puzzletron checkpoint, else a no-op."""
if not _ANYMODEL_AVAILABLE or not pretrained:
return contextlib.nullcontext()
try:
descriptor = mtpz.anymodel.resolve_descriptor_from_pretrained(
pretrained, trust_remote_code=trust_remote_code
)
except (ValueError, AttributeError):
return contextlib.nullcontext()
return mtpz.anymodel.deci_x_patcher(model_descriptor=descriptor)
def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict | None = None) -> T:
"""Overrides the HFLM.create_from_arg_obj"""
"""Override HFLM.create_from_arg_obj to add quantization, sparsity, and Puzzletron support."""
quant_cfg = arg_dict.pop("quant_cfg", None)
auto_quantize_bits = arg_dict.pop("auto_quantize_bits", None)
@@ -72,7 +100,10 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
# Enable automatic save/load of modelopt state huggingface checkpointing
mto.enable_huggingface_checkpointing()
model_obj = cls(**arg_dict, **additional_config)
with _anymodel_patcher_context(
arg_dict.get("pretrained"), arg_dict.get("trust_remote_code", False)
):
model_obj = cls(**arg_dict, **additional_config)
model_obj.tokenizer.padding_side = "left"
if is_quantized(model_obj.model):
# return if model is already quantized
@@ -109,10 +140,28 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
return model_obj
def create_from_arg_string(
cls: type[T], arg_string: str, additional_config: dict | None = None
) -> T:
"""Override HFLM.create_from_arg_string to support Puzzletron checkpoints."""
args = utils.simple_parse_args_string(arg_string)
additional_config = {} if additional_config is None else additional_config
args2 = {k: v for k, v in additional_config.items() if v is not None}
mto.enable_huggingface_checkpointing()
with _anymodel_patcher_context(args.get("pretrained"), args.get("trust_remote_code", False)):
model_obj = cls(**args, **args2)
return model_obj
HFLM.create_from_arg_obj = classmethod(create_from_arg_obj)
HFLM.create_from_arg_string = classmethod(create_from_arg_string)
def setup_parser_with_modelopt_args():
"""Extend the lm-eval argument parser with ModelOpt quantization and sparsity options."""
parser = setup_parser()
parser.add_argument(
"--quant_cfg",
examples/megatron_bridge/README.md
+24 -4
View File
@@ -46,6 +46,9 @@ Note that the default dataset for pruning and quantization is [`nemotron-post-tr
hf auth login --token <your token>
```
> [!WARNING]
> Use `python -m pip` instead of `pip` to avoid conflicts with the system-wide installed packages in the NeMo containers.
## Pruning
This section shows how to prune a HuggingFace model using Minitron algorithm in Megatron-Bridge framework. Checkout other available pruning algorithms, supported frameworks and models, and general pruning getting-started in the [pruning README](../pruning/README.md).
@@ -92,7 +95,7 @@ This section shows how to distill a student model from a teacher model in the Me
This can be used stand-alone or after [Pruning](#pruning) / [Post-Training Quantization](#post-training-quantization) to recover accuracy of the model by distilling from the original model (teacher).
The [distill.py](distill.py) script loads student and teacher models from HuggingFace checkpoints and saves the distilled model to `<output_dir>/checkpoints` in Megatron distributed checkpoint format.
The [distill.py](distill.py) script supports both standard HuggingFace checkpoints and [Puzzletron AnyModel](../puzzletron/README.md) checkpoints as student/teacher inputs. Just pass the checkpoint path via `--student_hf_path` / `--teacher_hf_path`. The distilled model is saved to `<output_dir>/checkpoints` in Megatron distributed checkpoint format.
### Data Preparation
@@ -158,9 +161,22 @@ torchrun --nproc_per_node 8 distill.py \
To run the distillation script on a Slurm cluster for multi-node training, you just need use `python` instead of `torchrun` and set the number of nodes using `#SBATCH --nodes=<num_nodes>` clause in your Slurm script.
### Convert Megatron checkpoint to Hugging Face format
### Converting to Hugging Face format (optional)
To convert the Megatron checkpoint from last iteration (or any intermediate iteration) to Hugging Face format, you need the pruned model config (`--output_hf_path` from `prune_minitron.py` script) and the distilled megatron checkpoint dir (`<distill_output_dir>/checkpoints/iter_<iter_number>`) to run the following command:
The distilled checkpoint is saved in Megatron distributed format. If you need a HuggingFace checkpoint, there are two ways to convert it:
**Inline** -- add `--hf_export_path` and `--student_hf_model` to the `distill.py` command to automatically convert the final checkpoint after distillation:
```bash
torchrun --nnodes 1 --nproc_per_node 8 distill.py \
... \
--hf_export_path /path/to/save/distilled_hf_ckpt \
--student_hf_model Qwen/Qwen3-4B
```
`--student_hf_model` should match the base architecture of the student (used as a template for export). For non-Puzzletron (i.e. standard) models, it should be same as `--student_hf_path`.
**Separate conversion** -- convert any saved iteration using the Megatron-Bridge conversion script:
```bash
uv run python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \
@@ -169,7 +185,11 @@ uv run python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py ex
--hf-path <path_to_save_distilled_hf_ckpt>
```
For more details, you can refer to the checkpoint conversion scripts in the [Megatron-Bridge README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion).
For more details, see the [Megatron-Bridge conversion README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion).
### Distillation Results
See [results/puzzletron.md](results/puzzletron.md) for MMLU results demonstrating knowledge distillation on Puzzletron-compressed student models.
## Post-Training Quantization
examples/megatron_bridge/distill.py
+92 -3
View File
@@ -15,17 +15,22 @@
"""Distillation script for Megatron-Bridge.
Loads student and teacher models directly from HuggingFace checkpoints (local or remote) and saves the distilled model
to `<output_dir>/checkpoints` in megatron distributed checkpoint format.
to `<output_dir>/checkpoints` in megatron distributed checkpoint or HuggingFace format.
See `README.md` in this directory for example usage and data preparation instructions.
"""
import argparse
import contextlib
import os
from dataclasses import fields
import torch
from megatron.bridge import AutoBridge
from megatron.bridge.models.distillation_provider import convert_to_distillation_provider
from megatron.bridge.models.distillation_provider import (
DistillationProvider,
convert_to_distillation_provider,
)
from megatron.bridge.recipes.utils.optimizer_utils import (
distributed_fused_adam_with_cosine_annealing,
)
@@ -43,13 +48,50 @@ from megatron.bridge.training.distill import distill
from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig
from megatron.core.datasets.utils import get_blend_from_list
from megatron.core.distributed import DistributedDataParallelConfig
from transformers import AutoConfig
import modelopt.torch.utils.distributed as dist
from modelopt.torch.utils import print_rank_0
with contextlib.suppress(ModuleNotFoundError):
import modelopt.torch.puzzletron.plugins.mbridge # noqa: F401
SEED = 1234
def _patched_to_cfg_dict(self):
"""Patched DistillationProvider.to_cfg_dict method for heterogeneous teacher and student models.
TODO: Upstream this patch to Megatron-Bridge.
"""
from megatron.bridge.training.utils.config_utils import _ConfigContainerBase
result = {"_target_": f"{self._super_class.__module__}.{self._super_class.__qualname__}"}
# Use fields from the actual student provider class, not DistillationProvider.
# DistillationProvider's __dataclass_fields__ only includes TransformerConfig fields
# (set at class definition time), missing GPTModelProvider-level fields like
# vocab_size, share_embeddings_and_output_weights, etc.
excluded_fields = {"teacher", "kd_config"}
for field in fields(self._super_class):
if field.name.startswith("_") or field.name in excluded_fields:
continue
if hasattr(self, field.name):
result[field.name] = _ConfigContainerBase._convert_value_to_dict(
getattr(self, field.name)
)
for field in fields(self):
if field.name.startswith("_") or field.name in excluded_fields:
continue
if field.name not in result:
result[field.name] = _ConfigContainerBase._convert_value_to_dict(
getattr(self, field.name)
)
return result
DistillationProvider.to_cfg_dict = _patched_to_cfg_dict
def get_args():
"""Parse command-line arguments."""
parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.")
@@ -124,12 +166,33 @@ def get_args():
)
parser.add_argument("--wandb_entity", type=str, help="Wandb entity name (optional)")
parser.add_argument("--wandb_exp_name", type=str, help="Wandb experiment name (optional)")
# Export arguments
parser.add_argument(
"--hf_export_path",
type=str,
default=None,
help=(
"Path where to save the HuggingFace export. "
"If provided, exports last iteration checkpoint to HF format after distillation."
),
)
parser.add_argument(
"--student_hf_model",
type=str,
required=False,
default=None,
help="HuggingFace model ID to use as template for export (e.g., Qwen/Qwen3-0.6B). "
"Should match the base architecture of the student model if --hf_export_path is provided.",
)
args = parser.parse_args()
# Sanity checks
if not args.use_mock_data and not args.data_paths:
raise ValueError("Must provide either --data_paths or set --use_mock_data.")
if args.hf_export_path and not args.student_hf_model:
raise ValueError("Must provide --student_hf_model if --hf_export_path is provided.")
print_rank_0("\n==================== Arguments ====================")
for k, v in args.__dict__.items():
print_rank_0(f"{k:<35} {v}")
@@ -252,9 +315,35 @@ def main(args: argparse.Namespace):
print_rank_0("\nStarting distillation...")
distill(config)
print_rank_0(
f"\nDistillation done! Saved checkpoint to {checkpoint_dir} in megatron distributed checkpoint format.\n"
f"\nDistillation done! Saved checkpoint to {checkpoint_dir}"
" in megatron distributed checkpoint format.\n"
)
if args.hf_export_path:
print_rank_0(f"Exporting final distilled ckpt to HF format to {args.hf_export_path}")
# Save rank before destroying process group (dist.rank() won't work after destruction)
is_rank_0 = dist.rank() == 0
# Destroy process group on all ranks -- export_ckpt will create its own temporary one.
# This prevents cleanup from hanging (cleanup tries to barrier, but rank 0 would be gone).
dist.cleanup()
if is_rank_0:
export_bridge = AutoBridge.from_hf_pretrained(
args.student_hf_model, trust_remote_code=args.trust_remote_code
)
# Copy weights and remote code
export_bridge.export_ckpt(
megatron_path=f"{checkpoint_dir}/iter_{args.train_iters:07d}",
hf_path=args.hf_export_path,
show_progress=True,
strict=True,
)
# Copy config.json from student_hf_path (handles both local paths and HF model IDs)
AutoConfig.from_pretrained(
args.student_hf_path, trust_remote_code=args.trust_remote_code
).save_pretrained(args.hf_export_path)
if __name__ == "__main__":
dist.setup()
examples/megatron_bridge/results/puzzletron.md
+42
View File
@@ -0,0 +1,42 @@
# Puzzletron Distillation Results
The following MMLU results demonstrate knowledge distillation on student models that were first compressed using [Puzzletron](../../puzzletron/README.md). The original (uncompressed) model serves as the teacher, and distillation recovers accuracy lost during compression.
## Qwen3-8B compressed to 80% of original
The student was created by compressing Qwen3-8B to 80% of its original size using Puzzletron.
| Model | MMLU | Humanities | Other | Social Sci | STEM |
|-------|------|------------|-------|------------|------|
| Student (before distillation) | 0.5910 | 0.5046 | 0.6363 | 0.6831 | 0.5855 |
| Student (after distillation) | 0.6921 | 0.5906 | 0.7316 | 0.7975 | 0.7016 |
| Teacher (original Qwen3-8B) | 0.7493 | 0.6648 | 0.7856 | 0.8385 | 0.7526 |
MMLU accuracy improved from 59.10% to 69.21% (+10.11 pp) after distillation with just 100 iterations on WikiText-103, recovering 64% of the gap to the teacher model.
## Llama-3.1-8B-Instruct compressed to 50% of original
The student was created by compressing Llama-3.1-8B-Instruct to 50% of its original size using Puzzletron.
| Model | MMLU | Humanities | Other | Social Sciences | STEM |
|-------|------|------------|-------|-----------------|------|
| Student (before distillation) | 0.2316 | 0.2462 | 0.2292 | 0.2250 | 0.2274 |
| Student (after distillation) | 0.2960 | 0.3146 | 0.3085 | 0.2925 | 0.2768 |
| Teacher (original Llama-3.1-8B-Instruct) | 0.6839 | 0.7231 | 0.7038 | 0.7667 | 0.5911 |
## Llama-3.1-8B-Instruct compressed to 69% of original (regression)
The student was created by compressing Llama-3.1-8B-Instruct to ~69% of its original size using Puzzletron. This example shows regression due to overfitting on the small WikiText-103 dataset (100 iterations). MMLU was evaluated on a subset of 100 samples per task:
| Model | MMLU | Humanities | Other | Social Sciences | STEM |
|-------|------|------------|-------|-----------------|------|
| Student (before distillation) | 0.6626 | 0.7069 | 0.6892 | 0.7525 | 0.5574 |
| Student (after distillation) | 0.6496 | 0.6862 | 0.6677 | 0.7433 | 0.5532 |
| Teacher (original Llama-3.1-8B-Instruct) | 0.6839 | 0.7231 | 0.7038 | 0.7667 | 0.5911 |
MMLU decreased from 66.26% to 64.96% (-1.30 pp) -- the model overfitted to WikiText-103. This highlights the importance of using larger, more diverse datasets for distillation.
## Recommendations
- **Use larger datasets** for production distillation (e.g., [Nemotron-Pretraining-SFT-v1](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-SFT-v1)) to avoid overfitting as shown in the regression case above.
- **Train for more iterations** to ensure proper convergence.
examples/pruning/README.md
+1
View File
@@ -7,6 +7,7 @@ Pruning can involve removal (prune) of Linear and Conv layers; and Transformer a
This section focuses on applying Model Optimizer's state-of-the-art complementary pruning modes to enable you to search for the best subnet architecture from your provided base model:
1. [Minitron](https://arxiv.org/pdf/2408.11796): A pruning method developed by NVIDIA Research for pruning GPT (and later extended to Mamba, MoE, and Hybrid Transformer Mamba) models in NVIDIA Megatron-LM (M-LM) or Megatron-Bridge (M-Bridge) framework. It uses the activation magnitudes to prune the embedding hidden size; mlp ffn hidden size; transformer attention heads; mamba heads and head dimension; MoE number of experts, ffn hidden size, and shared expert intermediate size; and number of layers of the model.
1. [Puzzletron](../puzzletron/README.md): An advanced pruning method by NVIDIA using Mixed Integer Programming (MIP) based NAS search algorithm.
1. FastNAS: A pruning method recommended for Computer Vision models. Given a pretrained model, FastNAS finds the subnet which maximizes the score function while meeting the given constraints.
1. GradNAS: A light-weight pruning method recommended for language models like Hugging Face BERT, GPT-J. It uses the gradient information to prune the model's linear layers and attention heads to meet the given constraints.
examples/puzzletron/GPTOSS.md
+14
View File
@@ -0,0 +1,14 @@
## GptOss
With this release Puzzle algorithm supports only experts removal for `Gpt-Oss`.
This model comes as a quantized checkpoint i.e. MoE experts matrices are quantized with _MXFP4_ format.
In the pruning steps puzzle utilizes decompressed model (back to BF16) for statistics and scores computation.
This means, during the conversion to puzzle format we decompress the model and store it as a BF16.
Once the pruning is done i.e. experts to be removed are identified and the process is finished, user may want to get back the _MXFP4_ format of the checkpoint.
To do so, there is an additional script, that takes the original and the pruned checkpoint and outputs pruned checkpoint in _MXFP4_ format.
```bash
python -m modelopt.torch.puzzletron.anymodel.models.gpt_oss.gpt_oss_pruned_to_mxfp4 --student-path /workspaces/any_model_gpt_oss/mip/puzzle_solutions/stats_num_params_18014757184/solutions--checkpoints/solution_0/ --original-path /workspaces/source_model_checkpoints/openai_gpt-oss-20b/ --output-path /workspaces/any_model_gpt_oss/mip/puzzle_solutions/stats_num_params_18014757184/solutions--checkpoints/mxfp4-ckpt/ --num-layers 24
```
examples/puzzletron/README.md
+306
View File
@@ -0,0 +1,306 @@
# Puzzletron Algorithm Tutorial
This tutorial demonstrates how to compress large language models using the puzzletron algorithm based on the [Puzzle paper](https://arxiv.org/abs/2411.19146).
The goal of the algorithm it to find the most optimal modifications to MLP and attention layers of the model, resulting in a heterogeneous model architecture.
The supported modifications are:
- `ffn_intermediate_size`: different FFN intermediate sizes
- `attention op/noop`: complete removal of attention layers
To use the Puzzle algorithm effectively, we need to specify the target number of parameters and/or the memory. The final stage is based on Mixed-Integer Programming (MIP) algorithm to find the most optimal combination of layer modifications that satisfy the target requirements.
In this example, we compress the [Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) model reducing GPU memory usage from 113 GiB to 96 GiB (15% reduction) with less than 1% regression in the token_accuracy_top_10 metric. Other supported models should be compressed in a similar way. For GptOss there is one [additional step to be performed](GPTOSS.md).
> **Note:** Other models are also supported. See the [configs](./configs/) directory for additional model configurations (e.g., Llama-3.2-3B-Instruct on 1x H100, Qwen2.5-7B-Instruct on 1x H100, Qwen3-8B on 1x H100, Nemotron-Nano-12B-v2 on 1x H100, Mistral-Small-24B-Instruct-2501 on 4x H100). For information on adding support for new models, see the [AnyModel Guide](../../modelopt/torch/puzzletron/anymodel/README.md).
## Environment
### Container setup (NeMo)
The recommended way to run puzzletron is inside an NVIDIA NeMo container (e.g. `nvcr.io/nvidia/nemo:26.02`). NeMo containers ship a pre-installed `nvidia-modelopt` that does not include the puzzletron extras so you need to replace it with an editable install from this repo.
> [!WARNING]
> Use `python -m pip` instead of `pip` to avoid conflicts with the system-wide installed packages in the NeMo containers.
> [!NOTE]
> NeMo containers ship `nvidia-lm-eval` which may conflict with `lm-eval` that is used for evaluation, hence we uninstall and replace it with `lm-eval` from the repo.
Once inside the container with the repo available, install dependencies from the repo root:
```bash
python -m pip uninstall nvidia-lm-eval -y 2>/dev/null
python -m pip install -e ".[hf,puzzletron,dev-test]"
python -m pip install -r examples/puzzletron/requirements.txt
```
To verify the install, you can run the GPU tests as a smoke check:
```bash
python -m pytest tests/gpu/torch/puzzletron/test_puzzletron.py -k "Qwen3-8B"
```
### Hardware
- For this example we are using 2x NVIDIA H100 80GB HBM3 to show multi-GPU steps. You can use also use a single GPU.
- To make use of [Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) and [Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2), you need to accept the terms and conditions for the corresponding model and the dataset in the Huggingface Hub. Log in to the Huggingface Hub and enter your HF token.
```bash
hf auth login --token <your token>
```
## Compress the Model
1. Download and prepare the [Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2).
dataset split: "code", "math", "stem", "chat", excluding reasoning samples (2.62GB)
```bash
python -m modelopt.torch.puzzletron.dataset.prepare_dataset \
--dataset_name nvidia/Nemotron-Post-Training-Dataset-v2 \
--output_dir path/to/Nemotron-Post-Training-Dataset-v2
```
2. Specify the `puzzle_dir`, `input_hf_model_path`, `dataset_path`, `intermediate_size_list`, and `target_memory` arguments in the [llama-3_1-8B_pruneffn_memory.yaml](./configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml) configuration file.
- `puzzle_dir` indicates a new directory for saving the resulting model.
- `input_hf_model_path` indicates the local directory with the input model checkpoint.
- `dataset_path` indicates the directory with the dataset downloaded earlier.
**_NOTE:_**
How to choose `intermediate_size_list`?
The list specifies the candidate FFN sizes that we wish to search over. It is recommended to choose several pruning sizes (e.g. 15%, 20%, 30% etc of the original). Note that the values must be hardware-friendly (divisible by a 256) to avoid issues with tensor operations in subsequent steps.
Let's first shoot for 32% GPU memory reduction setting `target_memory = 78_000` MiB. This means that the algorithm will choose the candidates with highest accuracy that also meet the specified requirements.
We can also set the target size of the resulting model using `num_params = 7_000_000_000`. This will be used as an upper bound for the number of parameters of the model.
3. Run the puzzletron pipeline.
```bash
torchrun --nproc_per_node 2 examples/puzzletron/main.py --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml 2>&1 | tee ./log.txt | grep "Puzzletron Progress"
```
This will save the full output to `log.txt` and display the following progress on screen:
```bash
[2025-11-02 12:06:34][rank-0][main.py:71] Puzzletron Progress 1/8: starting puzzletron pipeline
[2025-11-02 12:06:45][rank-0][puzzletron_nas_plugin.py:123] Puzzletron Progress 2/8: converting model from HF to DeciLM (single-gpu)
[2025-11-02 12:07:07][rank-0][puzzletron_nas_plugin.py:132] Puzzletron Progress 3/8: scoring pruning activations (multi-gpu)
[2025-11-02 12:11:36][rank-0][puzzletron_nas_plugin.py:137] Puzzletron Progress 4/8: pruning the model and saving pruned checkpoints (single-gpu)
[2025-11-02 12:12:20][rank-0][puzzletron_nas_plugin.py:217] Puzzletron Progress 5/8: building replacement library and subblock statistics (single-gpu)
[2025-11-02 12:12:21][rank-0][puzzletron_nas_plugin.py:222] Puzzletron Progress 6/8: calculating one block scores (multi-gpu)
[2025-11-02 12:50:41][rank-0][puzzletron_nas_plugin.py:226] Puzzletron Progress 7/8: running MIP and realizing models (multi-gpu)
[2025-11-02 12:52:34][rank-0][main.py:115] Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)
```
Once the process is complete, the resulting network architecture will be recorded in `log.txt` for your review:
```bash
...
block_0: attention gqa_4 ffn intermediate_14336
block_1: attention gqa_4 ffn intermediate_14336
block_2: attention gqa_4 ffn intermediate_14336
block_3: attention gqa_4 ffn intermediate_14336
block_4: attention gqa_4 ffn intermediate_14336
block_5: attention gqa_4 ffn intermediate_14336
block_6: attention gqa_4 ffn intermediate_14336
block_7: attention gqa_4 ffn intermediate_14336
block_8: attention gqa_4 ffn intermediate_14336
block_9: attention gqa_4 ffn intermediate_14336
block_10: attention gqa_4 ffn intermediate_14336
block_11: attention gqa_4 ffn intermediate_14336
block_12: attention gqa_4 ffn intermediate_14336
block_13: attention gqa_4 ffn intermediate_14336
block_14: attention gqa_4 ffn intermediate_14336
block_15: attention gqa_4 ffn intermediate_14336
block_16: attention gqa_4 ffn intermediate_14336
block_17: attention no_op ffn intermediate_14336
block_18: attention no_op ffn intermediate_14336
block_19: attention no_op ffn intermediate_14336
block_20: attention no_op ffn intermediate_14336
block_21: attention no_op ffn intermediate_14336
block_22: attention no_op ffn intermediate_14336
block_23: attention no_op ffn intermediate_14336
block_24: attention no_op ffn intermediate_14336
block_25: attention no_op ffn intermediate_14336
block_26: attention no_op ffn intermediate_14336
block_27: attention no_op ffn intermediate_14336
block_28: attention no_op ffn intermediate_14336
block_29: attention gqa_4 ffn intermediate_14336
block_30: attention gqa_4 ffn intermediate_14336
block_31: attention gqa_4 ffn intermediate_14336
[2025-11-02 04:53:11,332]^[[92m[rank-0]^[[0m[run_puzzle.py:295] Total costs: {'stats.memory_mib': 75796.4140625, 'stats.ffn_num_params': 5637275648, 'stats.num_kv_heads': 160, 'stats.kv_cache_memory_mib': 61440.0, 'stats.ffn_memory_mib': 10752.25, 'stats.attention_memory_mib': 63040.15625, 'stats.attention_num_params': 838942720, 'stats.num_params': 7526895616, 'stats.has_attention': 20, 'stats.has_ffn': 32}
...
################################################################
validate_model_and_extract_token_probs(model_name='teacher')
################################################################
...
Average losses = {'lm_loss': 1.118250765837729, 'token_accuracy_top_1': 0.7331905364990234, 'token_accuracy_top_5': 0.9094219207763672, 'token_accuracy_top_10': 0.9423646926879883}
...
################################################################
validate_model_with_kl_div(model_name='solution_0', is_calc_kl_div=True)
################################################################
....
Average losses = {'lm_loss': 1.7577573340386152, 'token_accuracy_top_1': 0.6225490570068359, 'token_accuracy_top_5': 0.846257209777832, 'token_accuracy_top_10': 0.8987817764282227}
```
30% GPU memory reduction leads to nearly 5% regression in token_accuracy_top_10 metric (0.898 / 0.942).
## Re-run MIP Search with different constraints
If you want to try different constraints without re-running the expensive pruning and scoring steps, use the `--mip-only` flag.
This assumes pruning, replacement library building, NAS scoring, and subblock stats calculation have already been completed.
For example, let's set `target_memory: 96_000` in `llama-3_1-8B_pruneffn_memory.yaml`.
```bash
torchrun --nproc_per_node 2 examples/puzzletron/main.py \
--config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \
--mip-only 2>&1 | tee ./log.txt | grep "Puzzletron Progress"
```
This will generate the following network architecture (see `log.txt`):
```bash
block_0: attention gqa_4 ffn intermediate_14336
block_1: attention gqa_4 ffn intermediate_14336
block_2: attention gqa_4 ffn intermediate_14336
block_3: attention gqa_4 ffn intermediate_14336
block_4: attention gqa_4 ffn intermediate_14336
block_5: attention gqa_4 ffn intermediate_14336
block_6: attention gqa_4 ffn intermediate_14336
block_7: attention gqa_4 ffn intermediate_14336
block_8: attention gqa_4 ffn intermediate_14336
block_9: attention gqa_4 ffn intermediate_14336
block_10: attention gqa_4 ffn intermediate_14336
block_11: attention gqa_4 ffn intermediate_14336
block_12: attention gqa_4 ffn intermediate_14336
block_13: attention gqa_4 ffn intermediate_14336
block_14: attention gqa_4 ffn intermediate_14336
block_15: attention gqa_4 ffn intermediate_14336
block_16: attention gqa_4 ffn intermediate_14336
block_17: attention gqa_4 ffn intermediate_14336
block_18: attention no_op ffn intermediate_14336
block_19: attention no_op ffn intermediate_14336
block_20: attention no_op ffn intermediate_14336
block_21: attention gqa_4 ffn intermediate_14336
block_22: attention no_op ffn intermediate_14336
block_23: attention no_op ffn intermediate_14336
block_24: attention no_op ffn intermediate_14336
block_25: attention gqa_4 ffn intermediate_14336
block_26: attention gqa_4 ffn intermediate_14336
block_27: attention gqa_4 ffn intermediate_14336
block_28: attention gqa_4 ffn intermediate_14336
block_29: attention gqa_4 ffn intermediate_14336
block_30: attention gqa_4 ffn intermediate_14336
block_31: attention gqa_4 ffn intermediate_14336
[2025-11-02 12:50:42,024]^[[92m[rank-0]^[[0m[run_puzzle.py:295] Total costs: {'stats.memory_mib': 94708.4609375, 'stats.has_ffn': 32, 'stats.ffn_memory_mib': 10752.25, 'stats.kv_cache_memory_mib': 79872.0, 'stats.attention_num_params': 1090625536, 'stats.ffn_num_params': 5637275648, 'stats.has_attention': 26, 'stats.num_params': 7778578432, 'stats.attention_memory_mib': 81952.203125, 'stats.num_kv_heads': 208}
...
################################################################
validate_model_with_kl_div(model_name='solution_0', is_calc_kl_div=True)
################################################################
Average losses = {'lm_loss': 1.2425934937782586, 'token_accuracy_top_1': 0.703862190246582, 'token_accuracy_top_5': 0.8954982757568359, 'token_accuracy_top_10': 0.9336576461791992
```
On the other hand, if you set `target_memory: 28_000`, you'll observe that the intermediate FFN sizes are significantly reduced in certain layers (see `log.txt` for details):
```bash
block_5: attention no_op ffn intermediate_11520
block_6: attention no_op ffn intermediate_14336
block_7: attention no_op ffn intermediate_8704
block_8: attention no_op ffn intermediate_14336
block_9: attention no_op ffn intermediate_3072
block_10: attention no_op ffn intermediate_11520
block_11: attention no_op ffn intermediate_11520
block_12: attention no_op ffn intermediate_11520
block_13: attention no_op ffn intermediate_11520
block_14: attention no_op ffn intermediate_3072
```
### MIP Sweep Mode
The **MIP sweep mode** lets you explore multiple memory compression rates in a single run and compare the accuracy-memory trade-offs.
#### Quick Start
1. Enable sweep in your config YAML (e.g., `llama-3_1-8B_pruneffn_memory.yaml`):
```yaml
mip:
sweep:
enabled: true
memory_compression_rates: [0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
output_csv: ${puzzle_dir}/mip_sweep_results.csv
```
2. Run the sweep:
```bash
torchrun --nproc_per_node 2 examples/puzzletron/main.py \
--config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \
--mip-only 2>&1 | tee ./log.txt | grep "Puzzletron Progress"
```
3. View results: The CSV file contains compression rates, memory usage, and accuracy metrics for each configuration.
#### Example Results
<img src="mip_sweep_example.png" alt="MIP Sweep Results" width="600">
The plot shows how token accuracy changes with different compression rates. Higher compression (0.5 = 50% of original memory) reduces accuracy, while lower compression maintains accuracy closer to the teacher model.
## Evaluation
Evaluate AnyModel checkpoints using [lm-eval](https://github.com/EleutherAI/lm-evaluation-harness) directly.
```bash
python examples/llm_eval/lm_eval_hf.py \
--model hf \
--model_args pretrained=path/to/checkpoint,dtype=bfloat16,parallelize=True \
--tasks mmlu \
--num_fewshot 5 \
--batch_size 4
```
For a quick smoke test, add `--limit 10`.
> **Alternative:** For server-based evaluation via an OpenAI-compatible endpoint,
> see [evaluation/nemo_evaluator_instructions.md](./evaluation/nemo_evaluator_instructions.md).
## Inference Performance Benchmarking
Now let's evaluate how much speedup we get with the compressed model in terms of throughput and latency.
- Install [vLLM from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source).
- Rearrange the model safetensors to be used for vLLM.
```bash
cd path/to/model
mv subblocks_safetensors/* .
sed -i 's+subblocks_safetensors/++g' model.safetensors.index.json
```
- Benchmark latency
```bash
vllm bench latency --model path/to/model --load-format safetensors
```
- Benchmark throughput
```bash
vllm bench throughput --model path/to/model --input-len 2000 --output-len 100 --load-format safetensors
```
## Knowledge Distillation
To recover degradation in the quality of the compressed model, we can use knowledge distillation. This allows transferring the capabilities of the original model to the pruned one.
See [Megatron-Bridge distillation](../megatron_bridge/README.md#distillation) for instructions on using Megatron-Bridge for knowledge distillation. The distillation script supports both standard HuggingFace and Puzzletron AnyModel checkpoints.
## Advanced Usage
Modify `llama-3_1-8B_pruneffn_memory.yaml` file for advanced compression scenarios.
examples/puzzletron/configs/gptoss-20b_remove_experts_memory/gptoss-20b.yaml
+110
View File
@@ -0,0 +1,110 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: gpt_oss
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
runtime_stats:
backend: trt_torch
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 45_000
num_params: 3_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/gptoss-20b_remove_experts_memory/gptoss-20b_remove_experts_memory.yaml
+17
View File
@@ -0,0 +1,17 @@
defaults:
- gptoss-20b
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/openai/gpt-oss-20b
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for compression outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 16_000 # 45 GiB
examples/puzzletron/configs/gptoss-20b_remove_experts_memory/pruning/ffn_pruning.yaml
+21
View File
@@ -0,0 +1,21 @@
defaults:
- pruning_defaults
eval_samples: 2500 #10
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/expert_removal/${pruning.experiment_id}
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.expert_removal_pruning_mixin.ExpertRemovalPruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.gpt_oss.gpt_oss_model_descriptor.GptOssExpertRemovalLayerDescriptor
target_name: "mlp.router"
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.expert_removal_hooks.RankedChoiceVotingHook}
activation_hooks_kwargs: # Additional kwargs to pass to the hook init
num_experts_to_keep_list: [24, 16, 8] # num_experts in teacher is 128
mlp_init_mode: "ExpertRemoval"
mlp_init_config_yaml:
expert_scores_key: "expert_ranks"
layer_prefix_template: "model.layers.{layer_idx}.mlp.router"
examples/puzzletron/configs/gptoss-20b_remove_experts_memory/pruning/pruning_defaults.yaml
+34
View File
@@ -0,0 +1,34 @@
defaults:
- /validate_model_defaults
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
descriptor: ${descriptor}
# Data:
eval_samples: 10_000
micro_batch_size: 1
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate" # PruneByActivationsLog
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_model_defaults.yaml
+18
View File
@@ -0,0 +1,18 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_solutions_defaults.yaml
+11
View File
@@ -0,0 +1,11 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/Llama-3_1-8B.yaml
+109
View File
@@ -0,0 +1,109 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: llama
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # ppath to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
runtime_stats:
backend: trt_torch
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 78_000
num_params: 7_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml
+26
View File
@@ -0,0 +1,26 @@
defaults:
- Llama-3_1-8B
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/meta-llama/Llama-3.1-8B-Instruct
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for puzzletron outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 78_000 # 78 GiB
# Memory sweep configuration (optional)
sweep:
enabled: false
memory_compression_rates: [0.5, 0.6, 0.7, 0.8, 0.9]
output_csv: ${puzzle_dir}/mip_sweep_results.csv
# FFN intermediate sizes to search over (heterogeneous architecture)
pruning:
intermediate_size_list: [3072, 5888, 8704, 11520] # teacher_intermediate_size is 14336
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml
+16
View File
@@ -0,0 +1,16 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: independent_kv_head_contribution
optimize_for: memory # IndependentKvHeadContributionHook implementation that consumes less memory
target_layer: "self_attn.o_proj"
layer_input_descriptors_path:
# n_heads_in_group: 4
# num_attention_heads: 32 # num query heads
# num_kv_heads: 32 / 4 = 8 # num_query_heads // n_heads_in_group
n_heads_in_group_list: [8, 16, 32] # num_kv_heads = [4, 2, 1]
gqa_init_mode: "PruneKVHeads"
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/ffn_pruning.yaml
+19
View File
@@ -0,0 +1,19 @@
defaults:
- pruning_defaults
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.llama.llama_model_descriptor.LlamaFFNIntermediateLayerDescriptor
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook}
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: iterative
target_layer: "mlp.down_proj"
layer_input_descriptors_path:
intermediate_size_list: [3072, 5888, 8704, 11520] # teacher_intermediate_size is 14336
mlp_init_mode: "PruneByActivationsLog"
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/hidden_dim_pruning.yaml
+15
View File
@@ -0,0 +1,15 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: layer_norm_contribution
target_layer: "layernorm"
# Hidden dimension pruning specific settings
hidden_size_list: [3072, 2048] # Target hidden sizes to prune to
hidden_size_init_mode: "PruneByChannelRanking"
mlp_init_mode: "Truncate" # TODO, make it work with CopyAsIs/FromTeacher
gqa_init_mode: "AverageKV" # TODO, make it work with CopyAsIs/FromTeacher
linear_init_mode: "FromTeacher"
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/pruning_defaults.yaml
+33
View File
@@ -0,0 +1,33 @@
defaults:
- /validate_model_defaults
descriptor: ${descriptor}
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
# Data:
eval_samples: 1000 # default is 10000
micro_batch_size: 4
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate" # PruneByActivationsLog
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml
+17
View File
@@ -0,0 +1,17 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_solutions_defaults.yaml
+10
View File
@@ -0,0 +1,10 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/Llama-3_2-3B.yaml
+110
View File
@@ -0,0 +1,110 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: llama
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
runtime_stats:
backend: trt_torch
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 45_000
num_params: 3_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/llama-3_2-3B_pruneffn_memory.yaml
+22
View File
@@ -0,0 +1,22 @@
defaults:
- Llama-3_2-3B
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/meta-llama/Llama-3.2-3B-Instruct
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for compression outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 45_000 # 45 GiB
# FFN intermediate sizes to search over (heterogeneous architecture)
# teacher_intermediate_size is 8192, so we use proportionally smaller values
pruning:
intermediate_size_list: [2048, 4096, 6144]
examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/pruning/ffn_pruning.yaml
+21
View File
@@ -0,0 +1,21 @@
defaults:
- pruning_defaults
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.llama.llama_model_descriptor.LlamaFFNIntermediateLayerDescriptor
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook}
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: iterative
target_layer: "mlp.down_proj"
layer_input_descriptors_path:
# Llama-3.2-3B has intermediate_size=8192, so we use proportionally smaller pruning sizes
intermediate_size_list: [2048, 4096, 6144]
mlp_init_mode: "PruneByActivationsLog"
examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/pruning/pruning_defaults.yaml
+33
View File
@@ -0,0 +1,33 @@
defaults:
- /validate_model_defaults
descriptor: ${descriptor}
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
# Data:
eval_samples: 1000 # default is 10000
micro_batch_size: 4
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate" # PruneByActivationsLog
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_model_defaults.yaml
+18
View File
@@ -0,0 +1,18 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_solutions_defaults.yaml
+11
View File
@@ -0,0 +1,11 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/Mistral-Small-24B.yaml
+109
View File
@@ -0,0 +1,109 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: mistral_small
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
runtime_stats:
backend: trt_torch
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 78_000
num_params: 24_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/mistral-small-24b-instruct-2501_pruneffn_memory.yaml
+21
View File
@@ -0,0 +1,21 @@
defaults:
- Mistral-Small-24B
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/mistralai/Mistral-Small-24B-Instruct-2501
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for compression outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 234_000 # 234 GiB
# FFN intermediate sizes to search over (heterogeneous architecture)
pruning:
intermediate_size_list: [8192, 16384, 24576] # teacher_intermediate_size is 32768
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/pruning/attn_pruning.yaml
+17
View File
@@ -0,0 +1,17 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: independent_kv_head_contribution
optimize_for: memory # IndependentKvHeadContributionHook implementation that consumes less memory
target_layer: "self_attn.o_proj"
layer_input_descriptors_path:
# Mistral Small 24B: 32 query heads, 8 KV heads
# n_heads_in_group = num_query_heads / num_kv_heads
# num_kv_heads = num_query_heads / n_heads_in_group
# Base: n_heads_in_group = 4, num_kv_heads = 8
n_heads_in_group_list: [8, 16, 32] # num_kv_heads = [4, 2, 1]
gqa_init_mode: "PruneKVHeads"
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/pruning/ffn_pruning.yaml
+20
View File
@@ -0,0 +1,20 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.mistral_small.mistral_small_model_descriptor.MistralFFNIntermediateLayerDescriptor
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook}
activation_hooks_kwargs:
method: iterative
target_layer: "mlp.down_proj"
layer_input_descriptors_path:
# FFN intermediate sizes to search over (heterogeneous architecture)
# teacher_intermediate_size is 32768
intermediate_size_list: [8192, 16384, 24576]
mlp_init_mode: "PruneByActivationsLog"
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/pruning/hidden_dim_pruning.yaml
+16
View File
@@ -0,0 +1,16 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: layer_norm_contribution
target_layer: "layernorm"
# Hidden dimension pruning specific settings
# Mistral Small 24B: hidden_size is 5120
hidden_size_list: [3072, 4096] # Target hidden sizes to prune to
hidden_size_init_mode: "PruneByChannelRanking"
mlp_init_mode: "Truncate" # TODO, make it work with CopyAsIs/FromTeacher
gqa_init_mode: "AverageKV" # TODO, make it work with CopyAsIs/FromTeacher
linear_init_mode: "FromTeacher"
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/pruning/pruning_defaults.yaml
+33
View File
@@ -0,0 +1,33 @@
defaults:
- /validate_model_defaults
descriptor: ${descriptor}
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
# Data:
eval_samples: 1000 # default is 10000
micro_batch_size: 4
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate" # PruneByActivationsLog
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_model_defaults.yaml
+17
View File
@@ -0,0 +1,17 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_solutions_defaults.yaml
+10
View File
@@ -0,0 +1,10 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/configs/nemotron-nano-12b-v2/nemotron_nano_12b_v2.yaml
+109
View File
@@ -0,0 +1,109 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: nemotron_h_v2
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
runtime_stats:
backend: trt_torch
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 90_000
num_params: 12_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/nemotron-nano-12b-v2/nemotron_nano_12b_v2_pruneffn_memory.yaml
+22
View File
@@ -0,0 +1,22 @@
defaults:
- nemotron_nano_12b_v2
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/nvidia/Nemotron-Nano-12B-v2
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for compression outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 90_000 # 90 GiB
# FFN intermediate sizes to search over (heterogeneous architecture)
# teacher_intermediate_size is 20480
pruning:
intermediate_size_list: [4352, 8448, 12544, 16384]
examples/puzzletron/configs/nemotron-nano-12b-v2/pruning/attn_pruning.yaml
+16
View File
@@ -0,0 +1,16 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: independent_kv_head_contribution
optimize_for: memory # IndependentKvHeadContributionHook implementation that consumes less memory
target_layer: "self_attn.o_proj"
layer_input_descriptors_path:
# n_heads_in_group: 4
# num_attention_heads: 32 # num query heads
# num_kv_heads: 32 / 4 = 8 # num_query_heads // n_heads_in_group
n_heads_in_group_list: [8, 16, 32] # num_kv_heads = [4, 2, 1]
gqa_init_mode: "PruneKVHeads"
examples/puzzletron/configs/nemotron-nano-12b-v2/pruning/ffn_pruning.yaml
+18
View File
@@ -0,0 +1,18 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.nemotron_h_v2.nemotron_h_v2_model_descriptor.NemotronHV2FFNIntermediateLayerDescriptor
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook}
activation_hooks_kwargs:
method: iterative
target_layer: "mixer.down_proj"
layer_input_descriptors_path:
intermediate_size_list: [256] # teacher_intermediate_size is 14336
mlp_init_mode: "PruneByActivationsLog"
examples/puzzletron/configs/nemotron-nano-12b-v2/pruning/hidden_dim_pruning.yaml
+15
View File
@@ -0,0 +1,15 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: layer_norm_contribution
target_layer: "layernorm"
# Hidden dimension pruning specific settings
hidden_size_list: [3072, 2048] # Target hidden sizes to prune to
hidden_size_init_mode: "PruneByChannelRanking"
mlp_init_mode: "Truncate" # TODO, make it work with CopyAsIs/FromTeacher
gqa_init_mode: "AverageKV" # TODO, make it work with CopyAsIs/FromTeacher
linear_init_mode: "FromTeacher"
examples/puzzletron/configs/nemotron-nano-12b-v2/pruning/pruning_defaults.yaml
+34
View File
@@ -0,0 +1,34 @@
defaults:
- /validate_model_defaults
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
descriptor: ${descriptor}
# Data:
eval_samples: 1000 # default is 10000
micro_batch_size: 4
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate"
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/nemotron-nano-12b-v2/validate_model_defaults.yaml
+17
View File
@@ -0,0 +1,17 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/nemotron-nano-12b-v2/validate_solutions_defaults.yaml
+10
View File
@@ -0,0 +1,10 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/pruning/attn_pruning.yaml
+16
View File
@@ -0,0 +1,16 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${modelopt.torch.puzzletron.pruning.activation_hooks_kwargs.method}/${modelopt.torch.puzzletron.pruning.experiment_id}
activation_hooks_kwargs:
method: independent_kv_head_contribution
optimize_for: memory # IndependentKvHeadContributionHook implementation that consumes less memory
target_layer: "self_attn.o_proj"
layer_input_descriptors_path:
# n_heads_in_group: 4
# num_attention_heads: 32 # num query heads
# num_kv_heads: 32 / 4 = 8 # num_query_heads // n_heads_in_group
n_heads_in_group_list: [8, 16, 32] # num_kv_heads = [4, 2, 1]
gqa_init_mode: "PruneKVHeads"
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/pruning/ffn_pruning.yaml
+18
View File
@@ -0,0 +1,18 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.qwen2.qwen2_model_descriptor.Qwen2FFNIntermediateLayerDescriptor
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook}
activation_hooks_kwargs:
method: iterative
target_layer: "mlp.down_proj"
layer_input_descriptors_path:
intermediate_size_list: [256] # teacher_intermediate_size is 14336
mlp_init_mode: "PruneByActivationsLog"
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/pruning/hidden_dim_pruning.yaml
+15
View File
@@ -0,0 +1,15 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${modelopt.torch.puzzletron.pruning.activation_hooks_kwargs.method}/${modelopt.torch.puzzletron.pruning.experiment_id}
activation_hooks_kwargs:
method: layer_norm_contribution
target_layer: "layernorm"
# Hidden dimension pruning specific settings
hidden_size_list: [3072, 2048] # Target hidden sizes to prune to
hidden_size_init_mode: "PruneByChannelRanking"
mlp_init_mode: "Truncate" # TODO, make it work with CopyAsIs/FromTeacher
gqa_init_mode: "AverageKV" # TODO, make it work with CopyAsIs/FromTeacher
linear_init_mode: "FromTeacher"
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/pruning/pruning_defaults.yaml
+34
View File
@@ -0,0 +1,34 @@
defaults:
- /validate_model_defaults
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
descriptor: ${descriptor}
# Data:
eval_samples: 1000 # default is 10000
micro_batch_size: 4
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate"
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/qwen2_5_7b_instruct.yaml
+109
View File
@@ -0,0 +1,109 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: qwen2
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
runtime_stats:
backend: trt_torch
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 78_000
num_params: 7_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/qwen2_5_7b_instruct_pruneffn_memory.yaml
+22
View File
@@ -0,0 +1,22 @@
defaults:
- qwen2_5_7b_instruct
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/Qwen/Qwen2.5-7B-Instruct
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for compression outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 78_000 # 78 GiB
# FFN intermediate sizes to search over (heterogeneous architecture)
# teacher_intermediate_size is 18944
pruning:
intermediate_size_list: [4096, 7808, 11520, 15104]
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_model_defaults.yaml
+17
View File
@@ -0,0 +1,17 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_solutions_defaults.yaml
+10
View File
@@ -0,0 +1,10 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/pruning/attn_pruning.yaml
+16
View File
@@ -0,0 +1,16 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: independent_kv_head_contribution
optimize_for: memory # IndependentKvHeadContributionHook implementation that consumes less memory
target_layer: "self_attn.o_proj"
layer_input_descriptors_path:
# n_heads_in_group: 4
# num_attention_heads: 32 # num query heads
# num_kv_heads: 32 / 4 = 8 # num_query_heads // n_heads_in_group
n_heads_in_group_list: [8, 16, 32] # num_kv_heads = [4, 2, 1]
gqa_init_mode: "PruneKVHeads"
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/pruning/ffn_pruning.yaml
+18
View File
@@ -0,0 +1,18 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
pruning_mixin:
_target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn
layer_descriptor:
_target_: modelopt.torch.puzzletron.anymodel.models.qwen3.qwen3_model_descriptor.Qwen3FFNIntermediateLayerDescriptor
hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook}
activation_hooks_kwargs:
method: iterative
target_layer: "mlp.down_proj"
layer_input_descriptors_path:
intermediate_size_list: [256] # teacher_intermediate_size is 14336
mlp_init_mode: "PruneByActivationsLog"
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/pruning/hidden_dim_pruning.yaml
+15
View File
@@ -0,0 +1,15 @@
defaults:
- pruning_defaults
activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id}
activation_hooks_kwargs:
method: layer_norm_contribution
target_layer: "layernorm"
# Hidden dimension pruning specific settings
hidden_size_list: [3072, 2048] # Target hidden sizes to prune to
hidden_size_init_mode: "PruneByChannelRanking"
mlp_init_mode: "Truncate" # TODO, make it work with CopyAsIs/FromTeacher
gqa_init_mode: "AverageKV" # TODO, make it work with CopyAsIs/FromTeacher
linear_init_mode: "FromTeacher"
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/pruning/pruning_defaults.yaml
+34
View File
@@ -0,0 +1,34 @@
defaults:
- /validate_model_defaults
model_name_or_path: ${teacher_dir}
experiment_id: ${pruning.eval_samples}samples_diverse_mini
activations_log_dir: ???
activation_hooks_kwargs: ???
descriptor: ${descriptor}
# Data:
eval_samples: 1000 # default is 10000
micro_batch_size: 4
dataset_path: ${dataset_path}
val_dataset_name: train
# Prune ckpts
pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id}
## FFN pruning
ffn_list:
mlp_init_mode: "Truncate"
## KV-heads pruning
n_heads_in_group_list:
gqa_init_mode: "AverageKV"
## Hidden dimension pruning
hidden_size_list:
hidden_size_init_mode: "PruneByChannelRanking"
linear_init_mode: "FromTeacher"
mlp_init_config_yaml:
activations_log_dir: ${pruning.activations_log_dir}
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/qwen3_8b.yaml
+109
View File
@@ -0,0 +1,109 @@
defaults:
- pruning: ffn_pruning
- scoring: ../validate_solutions_defaults
- realize_model: ../validate_solutions_defaults
- bypass:
- override hydra/hydra_logging: disabled
- _self_
puzzle_dir: ???
descriptor: qwen3
teacher_dir: ${puzzle_dir}/ckpts/teacher/
replacement_library_path: ${puzzle_dir}/replacement_library.json
dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2
skip_realize_model: false
build_replacement_library:
add_ffn_no_ops: true
add_attention_no_ops: true
calc_subblock_stats:
batch_sizes: [64, 96, 128]
prefill_seq_len: 4096
generation_seq_len: 4096
num_active_tokens_override: # Optional override for sequence lengths
prefill_queue_size: 0
allocate_prefill_query: false
runtime_stats:
backend: trt_torch
benchmark_iterations: # Set to a number (e.g., 1000) to enable runtime benchmarking
merge_with_existing_stats: false
subblock_stats_filename: "subblock_stats.json"
moe_stats_filename: "moe_stats.json"
scoring:
descriptor: ${descriptor}
solutions_to_validate:
skip_existing_solutions: true
replacement_library_path: ${replacement_library_path}
solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json}
teacher_dir: ${to_path:${teacher_dir}}
output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
mip:
single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}}
subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}}
output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions}
gathered_metrics_path:
puzzle_profile:
# puzzle_profile:
objective: metrics.cosine_embedding_loss_hidden_states
bigger_is_better: false
subblock_stats_args:
- batch_size: 96
weights_dtype: torch.bfloat16
activations_dtype: torch.bfloat16
kv_cache_dtype: torch.bfloat16
report_additional_costs:
- stats.memory_mib
- stats.num_params
- stats.num_kv_heads
- stats.has_attention
- stats.has_ffn
- stats.kv_cache_memory_mib
- stats.attention_memory_mib
- stats.ffn_memory_mib
- stats.ffn_num_params
- stats.attention_num_params
human_constraints:
target_memory: 78_000
num_params: 8_000_000_000
mip_constraints:
metric_overrides:
max_seconds_per_solution: 60
realize_model:
descriptor: ${descriptor}
teacher_dir: ${to_path:${teacher_dir}}
tokenizer_name: ${to_path:${teacher_dir}}
replacement_library_path: ${replacement_library_path}
save_models: true
solutions_path: # Filled dynamically
# Validate params
skip_validation: false # To enable validation of the model solution set `skip_validation` as False
eval_samples: 128
micro_batch_size: 1
seed: 42
shuffle_seed: 444
dataset_path: ${dataset_path}
nccl_timeout_minutes: ${timedelta_minutes:10}
# This section redirects Hydra outputs
hydra:
run:
dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S}
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/qwen3_8b_pruneffn_memory.yaml
+22
View File
@@ -0,0 +1,22 @@
defaults:
- qwen3_8b
- _self_
# Input Hugging Face model to compress
input_hf_model_path: /workspace/hf_models/Qwen/Qwen3-8B
# Dataset path for pruning and NAS scoring
dataset_path: /workspace/datasets/Nemotron-Post-Training-Dataset-v2
# Working directory for compression outputs
puzzle_dir: /workspace/puzzle_dir
# MIP memory constraint (in MiB)
mip:
human_constraints:
target_memory: 78_000 # 78 GiB
# FFN intermediate sizes to search over (heterogeneous architecture)
# teacher_intermediate_size is 12288
pruning:
intermediate_size_list: [2560, 5120, 7424, 9984]
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_model_defaults.yaml
+17
View File
@@ -0,0 +1,17 @@
model_dtype: torch.bfloat16 # dtype to cast the model for validate_model
autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model
block_size: 8192
bos_rate: 0.5
data_column: messages
val_dataset_name: valid
shuffle_seed: 81436
seed: 42
fim_rate: 0
fim_spm_rate: 0
source_datasets_to_discard:
varlen: false
write_results: false
calc_losses_on_cpu: false
activations_log_dir:
model_name_or_path:
load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn}
examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_solutions_defaults.yaml
+10
View File
@@ -0,0 +1,10 @@
defaults:
- /validate_model_defaults
- _self_
solutions_to_validate:
skip_validation: false
save_models: false
bigger_is_better: false
sort_solutions_by:
calculate_full_score_ablations: false
examples/puzzletron/evaluation/hf_deployable_anymodel.py
+720
View File
@@ -0,0 +1,720 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
import json
import logging
from typing import Any
import numpy as np
import torch
from nemo_deploy import ITritonDeployable
from nemo_deploy.utils import broadcast_list, cast_output, str_ndarray2list
from nemo_export_deploy_common.import_utils import (
MISSING_TRITON_MSG,
UnavailableError,
null_decorator,
)
from peft import PeftModel
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.puzzletron as mtpz
try:
from pytriton.decorators import batch
from pytriton.model_config import Tensor
HAVE_TRITON = True
except (ImportError, ModuleNotFoundError):
from unittest.mock import MagicMock
HAVE_TRITON = False
batch = MagicMock()
Tensor = MagicMock()
batch = null_decorator
LOGGER = logging.getLogger("NeMo")
SUPPORTED_TASKS = ["text-generation"]
class HuggingFaceLLMDeploy(ITritonDeployable):
"""A Triton inference server compatible wrapper for HuggingFace models.
This class provides a standardized interface for deploying HuggingFace models
in Triton inference server. It supports various NLP tasks and handles model
loading, inference, and deployment configurations.
Args:
hf_model_id_path (Optional[str]): Path to the HuggingFace model or model identifier.
Can be a local path or a model ID from HuggingFace Hub.
hf_peft_model_id_path (Optional[str]): Path to the PEFT model or model identifier.
Can be a local path or a model ID from HuggingFace Hub.
tokenizer_id_path (Optional[str]): Path to the tokenizer or tokenizer identifier.
If None, will use the same path as hf_model_id_path.
model (Optional[AutoModel]): Pre-loaded HuggingFace model.
tokenizer (Optional[AutoTokenizer]): Pre-loaded HuggingFace tokenizer.
tokenizer_padding (bool): Whether to enable padding in tokenizer. Defaults to True.
tokenizer_truncation (bool): Whether to enable truncation in tokenizer. Defaults to True.
tokenizer_padding_side (str): Which side to pad on ('left' or 'right'). Defaults to 'left'.
task (str): HuggingFace task type (e.g., "text-generation"). Defaults to "text-generation".
**hf_kwargs: Additional keyword arguments to pass to HuggingFace model loading.
"""
def __init__(
self,
hf_model_id_path: str | None = None,
hf_peft_model_id_path: str | None = None,
tokenizer_id_path: str | None = None,
model: AutoModel | None = None,
tokenizer: AutoTokenizer | None = None,
tokenizer_padding=True,
tokenizer_truncation=True,
tokenizer_padding_side="left",
task: str | None = "text-generation",
torch_dtype: torch.dtype | None = "auto",
device_map: str | None = "auto",
**hf_kwargs,
):
if not HAVE_TRITON:
raise UnavailableError(MISSING_TRITON_MSG)
if hf_model_id_path is None and model is None:
raise ValueError("hf_model_id_path or model parameters has to be passed.")
elif hf_model_id_path is not None and model is not None:
LOGGER.warning(
"hf_model_id_path will be ignored and the HuggingFace model set with model parameter will be used."
)
assert task in SUPPORTED_TASKS, "Task {} is not a support task.".format(task)
self.hf_model_id_path = hf_model_id_path
self.hf_peft_model_id_path = hf_peft_model_id_path
self.task = task
self.model = model
self.tokenizer = tokenizer
self.tokenizer_padding = tokenizer_padding
self.tokenizer_truncation = tokenizer_truncation
self.tokenizer_padding_side = tokenizer_padding_side
if tokenizer_id_path is None:
self.tokenizer_id_path = hf_model_id_path
else:
self.tokenizer_id_path = tokenizer_id_path
if model is None:
self._load(torch_dtype=torch_dtype, device_map=device_map, **hf_kwargs)
def _load(
self, torch_dtype: torch.dtype | None = "auto", device_map: str | None = "auto", **hf_kwargs
) -> None:
"""Load the HuggingFace pipeline with the specified model and task.
This method initializes the HuggingFace AutoModel classes using the provided model
configuration and task type. It handles the model and tokenizer loading
process.
Args:
torch_dtype (torch.dtype): Data type for the model. Defaults to "auto".
device_map (str): Device map for the model. Defaults to "auto".
**hf_kwargs: Additional keyword arguments to pass to the HuggingFace model loading.
Raises:
AssertionError: If task is not specified.
"""
assert self.task is not None, "A task has to be given for the generation task."
if self.task == "text-generation":
# =========================================================================
# BEGIN ANYMODEL PATCH
# Wraps model loading with deci_x_patcher for heterogeneous layer configs.
# See: modelopt/torch/puzzletron/anymodel/puzzformer/patcher.py
# =========================================================================
descriptor = mtpz.anymodel.resolve_descriptor_from_pretrained(
self.hf_model_id_path, trust_remote_code=hf_kwargs.get("trust_remote_code", False)
)
with mtpz.anymodel.deci_x_patcher(model_descriptor=descriptor):
self.model = AutoModelForCausalLM.from_pretrained(
self.hf_model_id_path,
torch_dtype=torch_dtype,
device_map=device_map,
**hf_kwargs,
)
# =========================================================================
# END ANYMODEL PATCH
# =========================================================================
if self.hf_peft_model_id_path is not None:
self.model = PeftModel.from_pretrained(self.model, self.hf_peft_model_id_path)
else:
raise ValueError("Task {} is not supported.".format(self.task))
num_gpus = torch.cuda.device_count()
# If there is only one GPU, move the model to GPU. If you are using device_map as "auto" or "balanced",
# the model will be moved to GPU automatically.
if device_map is None and num_gpus >= 1 and self.model.device.type != "cuda":
self.model.cuda()
self.tokenizer = AutoTokenizer.from_pretrained(
self.tokenizer_id_path,
trust_remote_code=hf_kwargs.pop("trust_remote_code", False),
padding=self.tokenizer_padding,
truncation=self.tokenizer_truncation,
padding_side=self.tokenizer_padding_side,
)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
def generate(
self,
**kwargs: Any,
) -> list[str]:
"""Generate text based on the provided input prompts.
This method processes input prompts through the loaded pipeline and
generates text according to the specified parameters.
Args:
**kwargs: Generation parameters including:
- text_inputs: List of input prompts
- max_length: Maximum number of tokens to generate
- num_return_sequences: Number of sequences to generate per prompt
- temperature: Sampling temperature
- top_k: Number of highest probability tokens to consider
- top_p: Cumulative probability threshold for token sampling
- do_sample: Whether to use sampling, default is False for greedy decoding
- echo: Whether to return prompt + generated text (True) or just generated text (False)
- return_full_text: Whether to return full text or only generated part
Returns:
If output logits and output scores are False:
List[str]: A list of generated texts, one for each input prompt.
If output logits and output scores are True:
Dict: A dictionary containing:
- sentences: List of generated texts
- logits: List of logits
- scores: List of scores
- input_lengths: List of input token lengths (for echo processing)
Raises:
RuntimeError: If the pipeline is not initialized.
"""
if not self.model:
raise RuntimeError("Model is not initialized")
inputs = self.tokenizer(
kwargs["text_inputs"],
return_tensors="pt",
padding=self.tokenizer_padding,
truncation=self.tokenizer_truncation,
)
# Store input lengths to extract only generated tokens later
input_lengths = [len(input_ids) for input_ids in inputs["input_ids"]]
# Get echo parameter (default False - only return generated text)
echo = kwargs.pop("echo", False)
kwargs.pop("text_inputs") # Remove text_inputs as it's already been tokenized
kwargs = {**inputs, **kwargs}
for key, val in kwargs.items():
if torch.is_tensor(val):
kwargs[key] = val.cuda()
with torch.no_grad():
generated_ids = self.model.generate(**kwargs)
return_dict_in_generate = kwargs.get("return_dict_in_generate", False)
if return_dict_in_generate:
# Handle dict output (when logits/scores are requested)
sequences = generated_ids["sequences"]
output = {"sentences": [], "input_lengths": input_lengths, "sequences": sequences}
if echo:
# Return full text (prompt + generated).
# HF model's generate returns the input/prompt tokens as well by default.
for i, seq in enumerate(sequences):
full_text = self.tokenizer.decode(seq, skip_special_tokens=True)
output["sentences"].append(full_text)
else:
# Extract only the generated tokens (skip input tokens).
# This is required as HF model's generate returns the input/prompt tokens
# as well by default. (return_full_text is specific to some models)
for i, seq in enumerate(sequences):
input_len = input_lengths[i] if i < len(input_lengths) else 0
generated_tokens = seq[input_len:] # Skip input tokens
generated_text = self.tokenizer.decode(
generated_tokens, skip_special_tokens=True
)
output["sentences"].append(generated_text)
if kwargs.get("output_logits", False):
output["logits"] = generated_ids["logits"]
if kwargs.get("output_scores", False):
output["scores"] = generated_ids["scores"]
else:
# Handle list output (normal case)
output = []
if echo:
# Return full text (prompt + generated), which is the default in case of HF model generate.
for i, seq in enumerate(generated_ids):
full_text = self.tokenizer.decode(seq, skip_special_tokens=True)
output.append(full_text)
else:
# Extract only the generated tokens (skip input tokens) as the default
# behavior returns the input/prompt tokens as well.
for i, seq in enumerate(generated_ids):
input_len = input_lengths[i] if i < len(input_lengths) else 0
generated_tokens = seq[input_len:] # Skip input tokens
generated_text = self.tokenizer.decode(
generated_tokens, skip_special_tokens=True
)
output.append(generated_text)
return output
def generate_other_ranks(self):
"""Generate function for ranks other than the rank 0."""
while True:
message = torch.empty(1, dtype=torch.long, device="cuda")
torch.distributed.broadcast(message, src=0)
if message == 0:
prompts = broadcast_list(data=[None], src=0)
(
temperature,
top_k,
top_p,
num_tokens_to_generate,
output_logits,
output_scores,
) = broadcast_list(data=[None], src=0)
return_dict_in_generate = False
if output_logits or output_scores:
return_dict_in_generate = True
self.generate(
text_inputs=prompts,
do_sample=False, # do_sample=False for greedy decoding
top_k=top_k,
top_p=top_p,
temperature=temperature,
max_new_tokens=num_tokens_to_generate,
output_logits=output_logits,
output_scores=output_scores,
return_dict_in_generate=return_dict_in_generate,
)
else:
return
@property
def get_triton_input(self):
inputs = (
Tensor(name="prompts", shape=(-1,), dtype=bytes),
Tensor(name="max_length", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="max_batch_size", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="top_k", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="top_p", shape=(-1,), dtype=np.single, optional=True),
Tensor(name="temperature", shape=(-1,), dtype=np.single, optional=True),
Tensor(name="random_seed", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="output_logits", shape=(-1,), dtype=np.bool_, optional=True),
Tensor(name="output_scores", shape=(-1,), dtype=np.bool_, optional=True),
)
return inputs
@property
def get_triton_output(self):
return (
Tensor(name="sentences", shape=(-1,), dtype=bytes),
Tensor(name="logits", shape=(-1,), dtype=np.single),
Tensor(name="scores", shape=(-1,), dtype=np.single),
)
@batch
def triton_infer_fn(self, **inputs: np.ndarray):
output_infer = {}
try:
prompts = str_ndarray2list(inputs.pop("prompts"))
temperature = inputs.pop("temperature")[0][0] if "temperature" in inputs else 1.0
top_k = int(inputs.pop("top_k")[0][0] if "top_k" in inputs else 1)
top_p = inputs.pop("top_p")[0][0] if "top_p" in inputs else 0
num_tokens_to_generate = (
inputs.pop("max_length")[0][0] if "max_length" in inputs else 256
)
output_logits = (
inputs.pop("output_logits")[0][0] if "output_logits" in inputs else False
)
output_scores = (
inputs.pop("output_scores")[0][0] if "output_scores" in inputs else False
)
return_dict_in_generate = False
if output_logits or output_scores:
return_dict_in_generate = True
if torch.distributed.is_initialized():
if torch.distributed.get_world_size() > 1:
torch.distributed.broadcast(
torch.tensor([0], dtype=torch.long, device="cuda"), src=0
)
broadcast_list(prompts, src=0)
broadcast_list(
data=[
temperature,
top_k,
top_p,
num_tokens_to_generate,
output_logits,
output_scores,
],
src=0,
)
output = self.generate(
text_inputs=prompts,
do_sample=False, # do_sample=False for greedy decoding
top_k=top_k,
top_p=top_p,
temperature=temperature,
max_new_tokens=num_tokens_to_generate,
output_logits=output_logits,
output_scores=output_scores,
return_dict_in_generate=return_dict_in_generate,
echo=False,
)
if isinstance(output, dict):
output_infer = {"sentences": cast_output(output["sentences"], np.bytes_)}
if "scores" in output:
output_scores = []
for r in output["scores"]:
lp = torch.tensor(r).cpu().detach().numpy()
if len(lp) == 0:
output_scores.append([0])
else:
output_scores.append(lp)
output_infer["scores"] = np.array(output_scores).transpose(1, 0, 2)
if "logits" in output:
output_logits = []
for r in output["logits"]:
lp = torch.tensor(r).cpu().detach().numpy()
if len(lp) == 0:
output_logits.append([0])
else:
output_logits.append(lp)
output_infer["logits"] = np.array(output_logits).transpose(1, 0, 2)
else:
output_infer = {"sentences": cast_output(output, np.bytes_)}
except Exception as error:
err_msg = "An error occurred: {}".format(str(error))
output_infer["sentences"] = cast_output([err_msg], np.bytes_)
return output_infer
def _compute_logprobs(
self,
prompts: list[str],
output_infer: dict[str, Any],
compute_logprob: bool,
n_top_logprobs: int,
echo: bool,
):
"""Compute log probabilities and top log probabilities from model scores.
Used by ray_infer_fn to provide OAI API compatible output for evaluations.
This method processes the raw scores from model generation to compute:
- Log probabilities for chosen tokens
- Top-k log probabilities for each position (if requested)
- Handles both prompt tokens (when echo=True) and generated tokens
Args:
prompts: List of input prompts
output_infer: Dictionary containing model outputs including scores, sequences, and input_lengths
compute_logprob: Whether to compute log probabilities
n_top_logprobs: Number of top log probabilities to return (0 to disable)
echo: Whether to include prompt token log probabilities
Returns:
Tuple[Optional[List], Optional[List]]:
- log_probs_list: List of log probabilities for each sample (None if not computed)
- top_logprobs_list: List of top-k log probabilities for each sample (None if not computed)
"""
# Tokenize the prompts to get prompt token IDs (needed for echo)
prompt_token_ids = None
prompt_inputs = None
if echo:
prompt_inputs = self.tokenizer(
prompts,
return_tensors="pt",
padding=self.tokenizer_padding,
truncation=self.tokenizer_truncation,
)
prompt_token_ids = prompt_inputs["input_ids"]
# Move to same device as model
for key, val in prompt_inputs.items():
if torch.is_tensor(val):
prompt_inputs[key] = val.cuda()
# Process each sample
log_probs_list = []
top_logprobs_list = []
for sample_idx in range(len(prompts)):
sample_log_probs = []
sample_top_logprobs = []
# Get the generated sequence for this sample
sequences = output_infer["sequences"][sample_idx]
# For echo, compute prompt token logprobs by running forward pass
if echo and prompt_token_ids is not None:
prompt_len = len(prompt_token_ids[sample_idx])
# Run forward pass on prompt to get logits for prompt tokens as scores in output_infer contains
# logits only for generated tokens.
with torch.no_grad():
# Create input for this specific sample
sample_prompt_input = {
key: val[sample_idx : sample_idx + 1] for key, val in prompt_inputs.items()
}
prompt_outputs = self.model(**sample_prompt_input)
prompt_logits = prompt_outputs.logits[0] # Shape: [seq_len, vocab_size]
# Calculate log probs for each prompt token (except the first BOS token)
for token_pos in range(1, prompt_len): # Start from 1 to skip BOS
# The logit at position i-1 predicts token at position i
logit_for_current_token = prompt_logits[token_pos - 1]
current_token_id = prompt_token_ids[sample_idx][token_pos].item()
# Calculate log probabilities
log_probs = torch.nn.functional.log_softmax(logit_for_current_token, dim=-1)
chosen_log_prob = log_probs[current_token_id].item()
sample_log_probs.append(chosen_log_prob)
# Calculate top log probabilities if requested
if n_top_logprobs > 0:
top_log_probs_dict = {}
top_k_values, top_k_indices = torch.topk(
log_probs, min(n_top_logprobs, len(log_probs))
)
for k_idx in range(len(top_k_indices)):
token_id = top_k_indices[k_idx].item()
token_str = self.tokenizer.decode([token_id])
top_log_probs_dict[token_str] = top_k_values[k_idx].item()
sample_top_logprobs.append(top_log_probs_dict)
# Process the scores for generated tokens
for token_idx, score_tensor in enumerate(output_infer["scores"]):
# Get the chosen token ID from the sequence
# Scores start after the prompt, so we need to offset
input_len = (
output_infer.get("input_lengths", [0])[sample_idx]
if "input_lengths" in output_infer
else 0
)
seq_idx = input_len + token_idx
if seq_idx < len(sequences):
chosen_token_id = (
sequences[seq_idx].item()
if hasattr(sequences[seq_idx], "item")
else sequences[seq_idx]
)
# Calculate log probabilities
log_probs = torch.nn.functional.log_softmax(score_tensor[sample_idx], dim=-1)
chosen_log_prob = log_probs[chosen_token_id].item()
sample_log_probs.append(chosen_log_prob)
# Calculate top log probabilities if requested
if n_top_logprobs > 0:
top_log_probs_dict = {}
top_k_values, top_k_indices = torch.topk(
log_probs, min(n_top_logprobs, len(log_probs))
)
for k_idx in range(len(top_k_indices)):
token_id = top_k_indices[k_idx].item()
token_str = self.tokenizer.decode([token_id])
top_log_probs_dict[token_str] = top_k_values[k_idx].item()
sample_top_logprobs.append(top_log_probs_dict)
log_probs_list.append(sample_log_probs)
if n_top_logprobs > 0:
top_logprobs_list.append(sample_top_logprobs)
# Return log probs and top logprobs
return_log_probs = log_probs_list if compute_logprob else None
return_top_logprobs = top_logprobs_list if n_top_logprobs > 0 else None
return return_log_probs, return_top_logprobs
def ray_infer_fn(self, inputs: dict[Any, Any]):
"""Perform inference using Ray with dictionary inputs and outputs.
Args:
inputs (Dict[Any, Any]): Dictionary containing input parameters:
- prompts: List of input prompts
- temperature: Sampling temperature (optional)
- top_k: Number of highest probability tokens to consider (optional)
- top_p: Cumulative probability threshold for token sampling (optional)
- max_tokens: Maximum number of tokens to generate (optional)
- compute_logprob: Whether to compute log probabilities (optional)
- n_top_logprobs: Number of top log probabilities to return (optional)
- echo: Whether to echo the prompt in output (optional)
Returns:
Dict[str, Any]: Dictionary containing:
- sentences: List of generated texts
- log_probs: Optional list of log probabilities if compute_logprob is True
- top_logprobs: Optional list of top log probabilities if n_top_logprobs > 0
"""
try:
prompts = inputs.pop("prompts")
temperature = inputs.pop("temperature", 1.0)
top_k = int(inputs.pop("top_k", 1))
top_p = inputs.pop("top_p", 0.0)
num_tokens_to_generate = inputs.pop("max_tokens", 256)
output_logits = inputs.pop("output_logits", False)
output_scores = inputs.pop("output_scores", False)
compute_logprob = inputs.pop("compute_logprob", False)
n_top_logprobs = inputs.pop("n_top_logprobs", 0)
echo = inputs.pop("echo", False)
output_infer = self._infer_fn_ray(
prompts=prompts,
temperature=temperature,
top_k=top_k,
top_p=top_p,
num_tokens_to_generate=num_tokens_to_generate,
output_logits=output_logits,
output_scores=output_scores,
compute_logprob=compute_logprob,
n_top_logprobs=n_top_logprobs,
echo=echo,
)
# Code to get logprobs (required in OAI API format for eval) from the scores in output_infer.
if (
(compute_logprob or n_top_logprobs > 0)
and "scores" in output_infer
and output_infer["scores"]
):
log_probs_list, top_logprobs_list = self._compute_logprobs(
prompts=prompts,
output_infer=output_infer,
compute_logprob=compute_logprob,
n_top_logprobs=n_top_logprobs,
echo=echo,
)
# Add to output
if log_probs_list is not None:
output_infer["log_probs"] = log_probs_list
if top_logprobs_list is not None:
# Convert to JSON strings for compatibility
output_infer["top_logprobs"] = [
json.dumps(top_logprobs) for top_logprobs in top_logprobs_list
]
# Remove raw outputs that are not needed in the final response
output_infer.pop("scores", None)
output_infer.pop("sequences", None)
output_infer.pop("input_lengths", None)
return output_infer
except Exception as error:
err_msg = "An error occurred: {}".format(str(error))
return {"sentences": [err_msg]}
def _infer_fn_ray(
self,
prompts,
temperature=1.0,
top_k=1,
top_p=0.0,
num_tokens_to_generate=256,
output_logits=False,
output_scores=False,
compute_logprob=False,
n_top_logprobs=0,
echo=False,
cast_output_func=None,
):
"""Common internal function for inference operations.
Args:
prompts: List of input prompts
temperature: Sampling temperature
top_k: Number of highest probability tokens to consider
top_p: Cumulative probability threshold for token sampling
num_tokens_to_generate: Maximum number of tokens to generate
output_logits: Whether to output logits
output_scores: Whether to output scores
compute_logprob: Whether to compute log probabilities
n_top_logprobs: Number of top log probabilities to return
echo: Whether to echo the prompt in output
cast_output_func: Optional function to cast output values
Returns:
Dict containing inference results with raw outputs
"""
# Enable return_dict if we need scores for logprobs or if output_logits/scores are requested
return_dict_in_generate = (
output_logits or output_scores or compute_logprob or n_top_logprobs > 0
)
# Enable output_scores if we need to compute logprobs. scores and logits from generate are both identical in
# case of greedy decoding. Hence setting output_scores to True when compute_logprob or n_top_logprobs > 0.
if compute_logprob or n_top_logprobs > 0:
output_scores = True
if torch.distributed.is_initialized():
if torch.distributed.get_world_size() > 1:
torch.distributed.broadcast(
torch.tensor([0], dtype=torch.long, device="cuda"), src=0
)
broadcast_list(prompts, src=0)
broadcast_list(
data=[
temperature,
top_k,
top_p,
num_tokens_to_generate,
output_logits,
output_scores,
],
src=0,
)
output = self.generate(
text_inputs=prompts,
do_sample=False, # do_sample=False for greedy decoding
top_k=top_k,
top_p=top_p,
temperature=temperature,
max_new_tokens=num_tokens_to_generate,
output_logits=output_logits,
output_scores=output_scores,
return_dict_in_generate=return_dict_in_generate,
echo=echo,
)
if isinstance(output, dict):
return output
else:
return {"sentences": output}
examples/puzzletron/evaluation/nemo_evaluator_instructions.md
+70
View File
@@ -0,0 +1,70 @@
# Evaluation with NeMo Evaluator (Alternative)
> **Recommended approach:** Use lm-eval for direct evaluation without a
> deployment server. See the main [README](../README.md#evaluation) for details.
Evaluate AnyModel checkpoints by deploying a local OpenAI-compatible completions endpoint and running benchmarks against it.
This flow requires Ray for serving the model and NeMo Export-Deploy (included in NeMo containers):
```bash
pip install -r examples/puzzletron/requirements.txt
```
**1. Deploy the model (2 GPUs example):**
We need to patch the `hf_deployable.py` script from Export-Deploy. Best way is to do it as a mount in docker run:
```bash
export MODELOPT_DIR=${PWD}/Model-Optimizer # or set to your local Model-Optimizer repository path if you have cloned it
if [ ! -d "${MODELOPT_DIR}" ]; then
git clone https://github.com/NVIDIA/Model-Optimizer.git ${MODELOPT_DIR}
fi
export DOCKER_IMAGE=nvcr.io/nvidia/nemo:26.02
docker run \
--gpus all \
--shm-size=16GB \
--net=host \
--ulimit memlock=-1 \
--rm -it \
-v ${MODELOPT_DIR}:/opt/Model-Optimizer \
-v ${MODELOPT_DIR}/modelopt:/opt/venv/lib/python3.12/site-packages/modelopt \
-v ${MODELOPT_DIR}/examples/puzzletron/evaluation/hf_deployable_anymodel.py:/opt/Export-Deploy/nemo_deploy/llm/hf_deployable.py \
-w /opt/Model-Optimizer/examples/megatron_bridge \
${DOCKER_IMAGE} bash
```
Alternatively you can manually update the file
```bash
# Install the AnyModel-patched deployable (first time only: backs up the original)
# /opt/Export-Deploy is the default path in NeMo containers — adjust if needed
cp /opt/Export-Deploy/nemo_deploy/llm/hf_deployable.py /opt/Export-Deploy/nemo_deploy/llm/hf_deployable.py.bak
cp examples/puzzletron/evaluation/hf_deployable_anymodel.py /opt/Export-Deploy/nemo_deploy/llm/hf_deployable.py
```
Now start ray server and deploy the model
```bash
# Start the server (blocks while running — use a separate terminal)
ray start --head --num-gpus 2 --port 6379 --disable-usage-stats
python /opt/Export-Deploy/scripts/deploy/nlp/deploy_ray_hf.py \
--model_path path/to/checkpoint \
--model_id anymodel-hf \
--num_gpus 2 --num_gpus_per_replica 2 --num_cpus_per_replica 16 \
--trust_remote_code --port 8083 --device_map "auto" --cuda_visible_devices "0,1"
```
**2. Run MMLU:**
```bash
eval-factory run_eval \
--eval_type mmlu \
--model_id anymodel-hf \
--model_type completions \
--model_url http://0.0.0.0:8083/v1/completions/ \
--output_dir examples/puzzletron/evals/mmlu_anymodel
```
For a quick debug run, add `--overrides "config.params.limit_samples=5"`.
examples/puzzletron/main.py
+170
View File
@@ -0,0 +1,170 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
"""
Main script for running the puzzletron algorithm on large language models (based on Puzzle paper https://arxiv.org/abs/2411.19146).
This script provides three modes:
1. Default mode: Runs the full puzzletron pipeline
2. MIP-only mode: Runs only the MIP search and realize models phase
3. MIP sweep mode: Runs MIP for multiple memory compression rates (enabled via config)
Usage:
# Full puzzletron pipeline
torchrun main.py --config ./configs/llama_3.2_1B_pruneffn_memory.yaml
# Only MIP search and realize models phase
torchrun main.py --config ./configs/llama_3.2_1B_pruneffn_memory.yaml --mip-only
# MIP sweep mode (set mip.sweep.enabled: true in config)
torchrun main.py --config ./configs/llama_3.2_1B_pruneffn_memory.yaml --mip-only
"""
import argparse
from datetime import timedelta
from pathlib import Path
import modelopt.torch.nas as mtn
import modelopt.torch.puzzletron as mtpz
import modelopt.torch.utils.distributed as dist
def parse_args():
"""Parse command line arguments."""
parser = argparse.ArgumentParser(
description="Compress large language models using the Puzzletron algorithm (based on Puzzle paper https://arxiv.org/abs/2411.19146)"
)
parser.add_argument(
"--config",
type=str,
required=True,
help="Path to the main config YAML file (e.g., ./configs/llama_3.2_1B_pruneffn_memory.yaml)",
)
parser.add_argument(
"--mip-only",
action="store_true",
help="Run only the MIP search and realize models phase (skip pruning and NAS scoring)",
)
return parser.parse_args()
def run_full_puzzletron(hydra_config_path: str):
"""Run the full puzzletron pipeline.
Args:
config_path: Path to the YAML configuration file
"""
mtpz.tools.mprint("Puzzletron Progress 1/8: starting puzzletron pipeline")
dist.setup(timeout=timedelta(minutes=10))
# Register Hydra custom resolvers (needed for config resolution)
mtpz.tools.register_hydra_resolvers()
hydra_config_path = Path(hydra_config_path).resolve()
hydra_config_dir = str(hydra_config_path.parent)
hydra_config_name = hydra_config_path.stem
# Load hydra config
hydra_cfg = mtpz.tools.initialize_hydra_config_for_dir(
config_dir=hydra_config_dir,
config_name=hydra_config_name,
overrides=[],
)
# Convert model (convert from HF to DeciLM, score pruning activations,
# prune the model and save pruned checkpoints)
input_model = mtpz.puzzletron_nas_plugin.PuzzletronModel()
converted_model = mtn.convert(
input_model,
mode=[
(
"puzzletron",
{
"puzzle_dir": str(hydra_cfg.puzzle_dir),
"input_model_path": hydra_cfg.input_hf_model_path,
"hydra_config_dir": hydra_config_dir,
"hydra_config_name": hydra_config_name,
"dataset_path": str(hydra_cfg.dataset_path),
},
)
],
)
# Run NAS search (build replacement library and compute stats,
# compute one block scores, run MIP and realize models)
mtn.search(
converted_model,
constraints={}, # this is not used as the search space is defined in the hydra config
dummy_input=None, # Not used
config={}, # this is not used as the search space is defined in the hydra config
)
dist.cleanup()
mtpz.tools.mprint("Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)")
def run_mip_only(hydra_config_path: str):
"""Run only the MIP search and realize models phase.
This assumes that pruning, replacement library building, NAS scoring, and subblock stats calculation
have already been completed.
Args:
hydra_config_path: Path to the YAML configuration file
"""
dist.setup(timeout=timedelta(minutes=10))
# Register Hydra custom resolvers (needed for config resolution)
mtpz.tools.register_hydra_resolvers()
hydra_config_path = Path(hydra_config_path).resolve()
hydra_config_dir = str(hydra_config_path.parent)
hydra_config_name = hydra_config_path.stem
# Load hydra config
hydra_cfg = mtpz.tools.initialize_hydra_config_for_dir(
config_dir=hydra_config_dir,
config_name=hydra_config_name,
overrides=[],
)
# Check if sweep mode is enabled
if hasattr(hydra_cfg.mip, "sweep") and hydra_cfg.mip.sweep.get("enabled", False):
mtpz.tools.mprint(
"Puzzletron Progress 7/8: running MIP sweep for multiple compression rates (multi-gpu)"
)
mtpz.mip.run_mip_sweep(hydra_cfg)
else:
# mip_and_realize_models (distributed processing)
# TODO: How to make it part of mnt.search() api, similarly to run_full_puzzletron() API
mtpz.tools.mprint("Puzzletron Progress 7/8: running MIP and realizing models (multi-gpu)")
mtpz.mip.launch_mip_and_realize_model(hydra_cfg)
dist.cleanup()
mtpz.tools.mprint("Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)")
def main():
args = parse_args()
if args.mip_only:
run_mip_only(hydra_config_path=args.config)
else:
run_full_puzzletron(hydra_config_path=args.config)
if __name__ == "__main__":
main()
examples/puzzletron/mip_sweep_example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 52 KiB

examples/puzzletron/requirements.txt
+5
View File
@@ -0,0 +1,5 @@
lm-eval==0.4.8
math-verify
ray
# Likely works for transformers v5 also, but we need to test it
transformers<5.0
modelopt/torch/export/plugins/__init__.py
+3
View File
@@ -25,3 +25,6 @@ from .vllm_fakequant_hf import *
with import_plugin("vllm_fakequant_megatron"):
from .vllm_fakequant_megatron import *
with import_plugin("hf_checkpoint_utils"):
from .hf_checkpoint_utils import *
modelopt/torch/export/plugins/hf_checkpoint_utils.py
+11 -13
View File
@@ -21,14 +21,13 @@ import shutil
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download, list_repo_files
from huggingface_hub import snapshot_download
from safetensors.torch import safe_open
from tqdm import tqdm
def copy_remote_code(
pretrained_model_path: str | os.PathLike,
save_directory: str | os.PathLike,
def copy_hf_ckpt_remote_code(
pretrained_model_path: str | os.PathLike, save_directory: str | os.PathLike
):
"""Copy remote code from pretrained model to save directory.
@@ -37,7 +36,7 @@ def copy_remote_code(
frameworks.
If ``pretrained_model_path`` is a local directory, Python files are copied directly.
If it is a HuggingFace Hub model ID, Python files are downloaded from the Hub first.
If it's a HF Hub model ID (e.g. ``nvidia/NVIDIA-Nemotron-Nano-12B-v2``), files are downloaded from the Hub.
Args:
pretrained_model_path: Local path to the pretrained model or HuggingFace Hub model ID.
@@ -45,18 +44,17 @@ def copy_remote_code(
"""
hf_checkpoint_path = Path(pretrained_model_path)
save_dir = Path(save_directory)
save_dir.mkdir(parents=True, exist_ok=True)
if hf_checkpoint_path.is_dir():
for py_file in hf_checkpoint_path.glob("*.py"):
if py_file.is_file():
shutil.copy(py_file, save_dir / py_file.name)
shutil.copy2(py_file, save_dir / py_file.name)
else:
# Hub model ID: download any top-level .py files (custom modeling code)
repo_id = str(pretrained_model_path)
for filename in list_repo_files(repo_id):
if "/" not in filename and filename.endswith(".py"):
local_path = hf_hub_download(repo_id=repo_id, filename=filename)
shutil.copy(local_path, save_dir / filename)
snapshot_download(
repo_id=str(pretrained_model_path),
local_dir=str(save_dir),
allow_patterns=["*.py"],
)
def load_multimodal_components(
modelopt/torch/export/unified_export_megatron.py
+2 -2
View File
@@ -45,7 +45,7 @@ from .model_config import (
QUANTIZATION_NONE,
QUANTIZATION_NVFP4,
)
from .plugins.hf_checkpoint_utils import copy_remote_code, load_multimodal_components
from .plugins.hf_checkpoint_utils import copy_hf_ckpt_remote_code, load_multimodal_components
from .plugins.mcore_common import all_mcore_hf_export_mapping
from .plugins.mcore_custom import (
CustomModuleMapping,
@@ -349,7 +349,7 @@ class GPTModelExporter:
torch.distributed.barrier()
if is_last_stage_main_rank and self._hf_config is not None:
copy_remote_code(pretrained_model_name_or_path, save_directory)
copy_hf_ckpt_remote_code(pretrained_model_name_or_path, save_directory)
# Newer versions of VLLM expect config.json with hf_quant_config
config_json_file = save_directory + "/config.json"
modelopt/torch/prune/importance_hooks/__init__.py
+1
View File
@@ -18,6 +18,7 @@ from modelopt.torch.utils import import_plugin
from .base_hooks import *
from .base_hooks_analysis import *
from .expert_removal_hooks import *
with import_plugin("megatron_hooks"):
from .plugins.megatron_hooks import *
modelopt/torch/prune/importance_hooks/base_hooks.py
+4 -3
View File
@@ -149,7 +149,8 @@ class ForwardHook(ABC):
torch.save(activations_log, activations_log_path)
if rank == 0:
args.activation_hooks_kwargs.pop("model")
if args.activation_hooks_kwargs is not None:
args.activation_hooks_kwargs.pop("model", None)
json_dump(OmegaConf.to_container(args, resolve=True), activations_log_dir / "args.json")
dist.barrier()
@@ -565,9 +566,9 @@ class IndependentKvHeadContributionHook(ForwardHook):
assert self.optimize_for in ["latency", "memory"]
self.hidden_size = model_config.hidden_size
self.n_heads_in_group = block_config.attention.n_heads_in_group
self.num_q_heads = model_config.num_attention_heads
self.num_kv_heads = self.num_q_heads // self.n_heads_in_group
self.num_kv_heads = block_config.attention.num_key_value_heads
self.n_heads_in_group = self.num_q_heads // self.num_kv_heads
self.head_dim = getattr(model_config, "head_dim", self.hidden_size // self.num_q_heads)
self.agg_kv_head_contributions = torch.zeros(
modelopt/torch/prune/importance_hooks/compare_module_outputs.py
+2 -1
View File
@@ -52,7 +52,8 @@ Step 2\: Compare outputs from different model variants\:
python compare_module_outputs.py \
--reference output_unpruned.pt \
--compare output_l2norm.pt \
--output-json comparison_stats.json
--output-json comparison_stats.json \
--no-weights-only
The saved file format\:
modelopt/torch/prune/importance_hooks/expert_removal_hooks.py
+404
View File
@@ -0,0 +1,404 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
"""MoE expert-removal and ranked-choice importance hooks (uses Puzzletron BlockConfig)."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
import torch
import transformers
from packaging.version import Version
from torch import nn
from .base_hooks import ForwardHook
if TYPE_CHECKING:
# Okay since this is only used for type hints else we should not import puzzletron here
# as its dependencies may not be installed
from modelopt.torch.puzzletron.block_config import BlockConfig
__all__ = [
"NemotronHRemoveExpertsIndependentHook",
"Qwen3VLRemoveExpertsIndependentHook",
"RankedChoiceVotingHook",
"RankedChoiceVotingHookNemotronH",
"RemoveExpertsIndependentHook",
]
class RemoveExpertsIndependentHook(ForwardHook, ABC):
"""Base hook for measuring expert importance in Mixture-of-Experts models.
This hook measures how much removing each expert affects the model output
by comparing outputs with and without each expert.
"""
def __init__(self, moe: nn.Module, activation_hooks_kwargs: dict):
"""Initialize the hook.
Args:
moe: The MoE module to analyze
activation_hooks_kwargs: Configuration dict containing block_config
"""
self.moe = moe
block_config: BlockConfig = activation_hooks_kwargs["block_config"]
self.num_local_experts = block_config.ffn.moe.num_local_experts
self.num_experts_per_tok = block_config.ffn.moe.num_experts_per_tok
# tensor of zeros of size num experts
self.diffs = ["mse", "cosine"]
some_param = next(self.moe.parameters())
self.diffs = {
k: torch.zeros(
size=(self.num_local_experts,), dtype=torch.float32, device=some_param.device
)
for k in self.diffs
}
self.call_count = 0
@abstractmethod
def get_router_logits_and_routed_experts(
self, hidden_states: torch.Tensor, router_logits: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
"""Extract router logits and expert outputs for measuring expert importance.
This method is called twice per forward pass:
1. First call (router_logits=None): Compute original routing and expert outputs
2. Second call (router_logits provided): Re-run with modified logits (expert disabled)
Args:
hidden_states: Input tensor of shape (batch, seq_len, hidden_dim)
router_logits: Optional pre-computed router logits. If None, compute from hidden_states.
Returns:
tuple of (router_logits, routed_experts):
- router_logits: Shape (num_tokens, num_local_experts)
- routed_experts: Shape (num_tokens, hidden_dim)
"""
raise NotImplementedError
def __call__(
self, module: nn.Module, args: tuple[torch.Tensor, ...], output: torch.Tensor
) -> None:
"""Forward hook that measures expert importance."""
hidden_states = args[0]
router_logits, original_routed_out = self.get_router_logits_and_routed_experts(
hidden_states
)
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
original_routed_out = original_routed_out.view(-1, original_routed_out.shape[-1])
_, router_indices = torch.topk(router_logits, self.num_experts_per_tok, dim=-1)
self.call_count += 1
for i_expert in range(self.num_local_experts):
expert_mask = router_indices == i_expert
is_token_routed_to_this_expert = expert_mask.any(dim=-1)
num_tokens_displaced = is_token_routed_to_this_expert.sum()
if num_tokens_displaced == 0:
continue
num_total_tokens = is_token_routed_to_this_expert.numel()
relevant_hidden_states = hidden_states[is_token_routed_to_this_expert, :]
router_logits_without_i = router_logits.clone()
router_logits_without_i[..., i_expert] = -float("inf") # disable expert i
router_logits_without_i = router_logits_without_i[is_token_routed_to_this_expert, :]
_, routed_out_without_i = self.get_router_logits_and_routed_experts(
relevant_hidden_states, router_logits_without_i
)
relevant_tokens_original_out = original_routed_out[is_token_routed_to_this_expert, :]
self.diffs["mse"][i_expert] += (
nn.functional.mse_loss(
relevant_tokens_original_out, routed_out_without_i, reduction="mean"
)
* num_tokens_displaced
/ num_total_tokens
)
self.diffs["cosine"][i_expert] += (
-nn.functional.cosine_similarity(
relevant_tokens_original_out, routed_out_without_i, dim=-1
).mean()
* num_tokens_displaced
/ num_total_tokens
)
def to_dict(self) -> dict[str, torch.Tensor]:
"""Convert accumulated statistics to dict format."""
expert_ranks_mse = torch.argsort(self.diffs["mse"])
expert_ranks_cosine = torch.argsort(self.diffs["cosine"])
return {
"expert_ranks_mse": expert_ranks_mse.cpu(),
"expert_ranks_cosine": expert_ranks_cosine.cpu(),
"cosine_diffs": (self.diffs["cosine"] / self.call_count).cpu(),
"mse_diffs": (self.diffs["mse"] / self.call_count).cpu(),
}
def accumulate(self) -> torch.Tensor:
"""Return accumulated expert importance scores."""
return self.diffs["mse"]
def state_dict(self) -> dict:
"""Return the internal state for checkpointing."""
return {
"diffs_mse": self.diffs["mse"].cpu(),
"diffs_cosine": self.diffs["cosine"].cpu(),
"call_count": self.call_count,
}
def load_state_dict(self, state_dict: dict) -> None:
"""Load the internal state from a checkpoint."""
self.diffs["mse"] = state_dict["diffs_mse"].to(self.diffs["mse"].device)
self.diffs["cosine"] = state_dict["diffs_cosine"].to(self.diffs["cosine"].device)
self.call_count = state_dict["call_count"]
class NemotronHRemoveExpertsIndependentHook(RemoveExpertsIndependentHook):
"""Expert removal importance hook for NemotronH models."""
def get_router_logits_and_routed_experts(
self, hidden_states: torch.Tensor, router_logits: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
"""Extract router logits and expert outputs for NemotronH MoE.
Based on NemotronHMOE forward, uses minimum ops to get router_logits and routed_experts.
"""
orig_shape = hidden_states.shape
# NemotronHMOE.gate forward, copied to extract router_logits
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
if router_logits is None:
router_logits = nn.functional.linear(
hidden_states.type(torch.float32), self.moe.gate.weight.type(torch.float32)
)
router_logits = router_logits.sigmoid()
router_logits = router_logits + self.moe.gate.e_score_correction_bias.unsqueeze(0)
topk_indices = self._get_topk_indices_without_correction_bias(router_logits)
topk_weights = router_logits.gather(1, topk_indices)
if self.moe.gate.norm_topk_prob:
denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
topk_weights /= denominator
topk_weights = topk_weights * self.moe.gate.routed_scaling_factor
# Routed experts forward
hidden_states = self.moe.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape)
return router_logits, hidden_states
@torch.no_grad()
def _get_topk_indices_without_correction_bias(self, scores: torch.Tensor) -> torch.Tensor:
"""Get topk indices without correction bias.
Same as NemotronHMOE.gate.get_topk_indices but without adding e_score_correction_bias.
"""
group_scores = (
scores.view(
-1, self.moe.gate.n_group, self.moe.gate.n_routed_experts // self.moe.gate.n_group
)
.topk(2, dim=-1)[0]
.sum(dim=-1)
)
group_idx = torch.topk(group_scores, k=self.moe.gate.topk_group, dim=-1, sorted=False)[1]
group_mask = torch.zeros_like(group_scores)
group_mask.scatter_(1, group_idx, 1)
score_mask = (
group_mask.unsqueeze(-1)
.expand(
-1, self.moe.gate.n_group, self.moe.gate.n_routed_experts // self.moe.gate.n_group
)
.reshape(-1, self.moe.gate.n_routed_experts)
)
scores_for_choice = scores.masked_fill(~score_mask.bool(), 0.0)
topk_indices = torch.topk(scores_for_choice, k=self.moe.gate.top_k, dim=-1, sorted=False)[1]
return topk_indices
class RankedChoiceVotingHook(ForwardHook):
"""Hook for ranking experts using ranked choice voting algorithm.
This hook tracks router decisions and uses ranked choice voting to determine
which experts are least important (can be pruned first).
"""
def __init__(self, router: nn.Module, activation_hooks_kwargs: dict):
"""Initialize the hook.
Args:
router: The router module (typically nn.Linear)
activation_hooks_kwargs: Configuration dict containing block_config
"""
self.router_argsort: list[torch.Tensor] = []
block_config: BlockConfig = activation_hooks_kwargs["block_config"]
self.top_k = block_config.ffn.moe.num_experts_per_tok
def __call__(
self, module: nn.Module, args: tuple[torch.Tensor, ...], output: torch.Tensor
) -> None:
"""Forward hook that records router decisions.
Args:
module: The router module
args: Tuple with one tensor entry (B, T, I)
output: Router logits of shape (B, T, E)
"""
router_logits = output[0] if isinstance(output, tuple) else output
num_experts = router_logits.shape[-1]
router_argsort = torch.argsort(router_logits, dim=-1, descending=True)
router_argsort = router_argsort.view(-1, num_experts).to(torch.int16).cpu()
self.router_argsort.append(router_argsort)
def to_dict(self) -> dict[str, torch.Tensor]:
"""Convert accumulated statistics to dict format using ranked choice voting."""
router_argsort = torch.concat(self.router_argsort, dim=0)
num_tokens, num_experts = router_argsort.shape
expert_ranks = torch.full((num_experts,), -1)
expert_counts_at_pruning_time = {}
expert_kept_per_iteration: list[list[int]] = []
expert_counts_per_iteration: list[dict[int, int]] = []
for rank in range(num_experts):
ids, counts = router_argsort[:, : self.top_k].unique(return_counts=True)
ids = ids.tolist()
counts = counts.tolist()
expert_counts = dict(zip(ids, counts))
expert_kept_per_iteration.append(ids)
expert_counts_per_iteration.append(expert_counts)
least_popular_expert, min_count = min(expert_counts.items(), key=lambda tup: tup[1])
expert_ranks[least_popular_expert] = rank
expert_counts_at_pruning_time[least_popular_expert] = min_count
router_argsort = router_argsort[router_argsort != least_popular_expert].view(
num_tokens, -1
)
zero_shot_expert_counts = torch.zeros((num_experts,), dtype=torch.long)
for expert_id, expert_counts_val in expert_counts_per_iteration[0].items():
zero_shot_expert_counts[expert_id] = expert_counts_val
# Compute zero-shot expert ranks (double argsort converts counts to rank positions)
zero_shot_expert_ranks = torch.argsort(torch.argsort(zero_shot_expert_counts))
return {
"expert_ranks": expert_ranks,
"zero_shot_expert_ranks": zero_shot_expert_ranks,
"expert_counts_at_pruning_time": expert_counts_at_pruning_time,
"expert_counts_per_iteration": expert_counts_per_iteration,
"top_k": self.top_k,
}
def accumulate(self) -> torch.Tensor:
"""Return accumulated expert ranks."""
if not self.router_argsort:
return torch.tensor([])
router_argsort = torch.concat(self.router_argsort, dim=0)
return router_argsort[:, 0].float()
def state_dict(self) -> dict:
"""Return the internal state for checkpointing."""
return {
"router_argsort": [tensor.cpu().clone() for tensor in self.router_argsort],
"top_k": self.top_k,
}
def load_state_dict(self, state_dict: dict) -> None:
"""Load the internal state from a checkpoint."""
self.router_argsort = [tensor.cpu() for tensor in state_dict["router_argsort"]]
self.top_k = state_dict["top_k"]
def get_progress_info(self) -> dict:
"""Get progress information."""
return {
"num_batches_processed": len(self.router_argsort),
"total_tokens_processed": sum(tensor.shape[0] for tensor in self.router_argsort)
if self.router_argsort
else 0,
}
class RankedChoiceVotingHookNemotronH(RankedChoiceVotingHook):
"""Ranked choice voting hook for NemotronH models.
In NemotronH, router_logits is an internal temporary state that never leaves
the forward() function. We reconstruct router_logits from the input hidden_states.
"""
def __call__(
self, module: nn.Module, args: tuple[torch.Tensor, ...], output: torch.Tensor
) -> None:
"""Forward hook that reconstructs router logits from hidden states."""
hidden_states = args[0]
hidden_states = hidden_states.view(-1, module.config.hidden_size)
router_logits = nn.functional.linear(
hidden_states.type(torch.float32), module.weight.type(torch.float32)
)
super().__call__(module, args, router_logits)
class Qwen3VLRemoveExpertsIndependentHook(RemoveExpertsIndependentHook):
"""Expert removal importance hook for Qwen3-VL models."""
def get_router_logits_and_routed_experts(
self, hidden_states: torch.Tensor, router_logits: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
"""Extract router logits and expert outputs for Qwen3-VL MoE.
Based on Qwen3VLMoeSparseMoe forward pass.
"""
orig_shape = hidden_states.shape
# Use hidden_states.shape[-1] instead of self.moe.hidden_size for transformers v5 compatibility
hidden_size = (
self.moe.hidden_size if hasattr(self.moe, "hidden_size") else hidden_states.shape[-1]
)
# Flatten to (num_tokens, hidden_size) for processing
hidden_states_flat = hidden_states.reshape(-1, hidden_size)
if router_logits is None:
router_logits = self.moe.gate(hidden_states_flat)
# In transformers vf the gate returns (logits, aux_loss) tuple
if isinstance(router_logits, tuple):
router_logits = router_logits[0]
routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
routing_weights, router_indices = torch.topk(
routing_weights, self.num_experts_per_tok, dim=-1
)
routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
routing_weights = routing_weights.to(hidden_states_flat.dtype)
if Version(transformers.__version__) >= Version("5.0"):
# transformers 5.x: grouped_mm_experts_forward expects
# (hidden_states_flat 2D, top_k_index, top_k_weights)
routed_out = self.moe.experts(hidden_states_flat, router_indices, routing_weights)
else:
# transformers 4.x: loop-based experts expects
# (hidden_states_3d 3D, routing_weights_full, router_indices)
batch_size = orig_shape[0] if hidden_states.ndim == 3 else 1
hidden_states_3d = hidden_states_flat.reshape(batch_size, -1, hidden_size)
router_weights = torch.zeros(
router_logits.shape, dtype=routing_weights.dtype, device=router_logits.device
).scatter_(1, router_indices, routing_weights)
routed_out = self.moe.experts(hidden_states_3d, router_weights, router_indices)
# Return in same shape as input
routed_out = routed_out.reshape(*orig_shape)
return router_logits, routed_out
modelopt/torch/puzzletron/README.md
+3
View File
@@ -0,0 +1,3 @@
Experimental model compression algorithm based on a Local Neural Architecture Search.
Based on the Puzzle paper: <https://arxiv.org/abs/2411.19146>
PoC for Llama 3.1 model.
modelopt/torch/puzzletron/__init__.py
+33
View File
@@ -0,0 +1,33 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# NOTE: Some modules also trigger factory registration as side effect
from . import (
activation_scoring,
anymodel,
block_config,
build_library_and_stats,
dataset,
entrypoint,
mip,
plugins,
pruning,
puzzletron_nas_plugin,
replacement_library,
scoring,
subblock_stats,
tools,
utils,
)
modelopt/torch/puzzletron/activation_scoring/__init__.py
+17
View File
@@ -0,0 +1,17 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
from .activation_hooks import *
from .score_pruning_activations import *
modelopt/torch/puzzletron/activation_scoring/activation_hooks/__init__.py
+16
View File
@@ -0,0 +1,16 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
from .utils import *
modelopt/torch/puzzletron/activation_scoring/activation_hooks/utils.py
+99
View File
@@ -0,0 +1,99 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
"""Provides a function to register activation hooks for a model.
Activation hooks are used to compute activation scores for pruning."""
from typing import Type
import torch
from modelopt.torch.prune.importance_hooks.base_hooks import ForwardHook as ActivationsHook
from ...tools.logger import aprint
from ...utils.dummy_modules import DummyBlock, DummyModule
__all__ = ["register_activation_hooks"]
def register_activation_hooks(
model,
activation_hooks_kwargs: dict,
pruning_mixin,
hook_class: Type[ActivationsHook],
) -> dict[str, ActivationsHook]:
"""Register activation hooks using the pruning mixin approach.
Args:
model: The model to register hooks on.
activation_hooks_kwargs: Keyword arguments passed to hook constructors.
pruning_mixin: The pruning mixin that defines which modules to hook.
hook_class: The hook class to instantiate for each module.
Returns:
Dictionary mapping module names to hook instances.
"""
activation_hooks_kwargs["model"] = model
if hook_class not in pruning_mixin.supported_hooks():
raise ValueError(
f"Hook class not supported for {pruning_mixin.__class__.__name__}, "
f"must be in {pruning_mixin.supported_hooks()}"
)
module_names_to_hook = pruning_mixin.get_module_names_to_hook(model)
activation_hooks = dict()
for block_idx, module_name in module_names_to_hook:
try:
module = model.get_submodule(module_name)
except AttributeError:
# Module doesn't exist on this rank's shard (e.g., in distributed setup)
continue
# Skip dummy modules - they don't have real activations to hook
if isinstance(module, (DummyModule, DummyBlock)):
continue
block_config = None
if block_idx is not None:
block_config = model.config.block_configs[block_idx]
curr_activation_hooks_kwargs = {
**activation_hooks_kwargs,
"block_config": block_config,
}
hook = hook_class(module, curr_activation_hooks_kwargs)
module.register_forward_hook(hook)
activation_hooks[module_name] = hook
if len(activation_hooks) == 0:
# In distributed mode, it's okay for a rank to have 0 hooks if it doesn't own
# the target modules (e.g., with hybrid patterns like "*-" where different
# ranks own different layer types). However, we still want to catch real bugs
# where no hooks are found at all.
is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
if is_distributed:
aprint(
"No hooks registered on this rank. This is expected if this rank "
"doesn't own any layers matching the hook pattern (e.g., in hybrid "
"patterns with distributed model sharding)."
)
else:
raise ValueError("couldn't find any hooks")
if len(activation_hooks) > 0:
aprint(f"Found the following hooks: {activation_hooks.keys()}")
return activation_hooks
modelopt/torch/puzzletron/activation_scoring/score_pruning_activations.py
+142
View File
@@ -0,0 +1,142 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
from pathlib import Path
import torch
from omegaconf import DictConfig
import modelopt.torch.utils.distributed as dist
from ..tools.logger import mprint
__all__ = ["launch_score_activations"]
def has_checkpoint_support(activation_hooks_kwargs: dict) -> bool:
"""Determine if the activation hook method has proper checkpoint support implemented.
Args:
activation_hooks_kwargs: Hook configuration
Returns:
bool: True if the hook method has save_state/load_state implemented
"""
method = activation_hooks_kwargs.get("method", "")
# Methods with implemented checkpoint support
supported_methods = {
"iterative", # IterativeChannelContributionHook: save_state/load_state implemented
"independent", # IndependentChannelContributionHook: save_state/load_state implemented
"stats", # RouterStatsHook: save_state/load_state implemented
"ranked_choice_voting", # RankedChoiceVotingHook: save_state/load_state implemented
}
return method in supported_methods
def check_scoring_completion(activations_log_dir: str, activation_hooks_kwargs=None) -> bool:
"""Check if scoring is already completed by looking for the expected output files.
Also checks if the scoring method is safe for resume.
Args:
activations_log_dir: Directory where activation logs should be stored
activation_hooks_kwargs: Hook configuration to check if resume is safe
Returns:
bool: True if scoring is completed (has rank files and args.json)
"""
# Only check completion on main process
if dist.is_master():
log_dir = Path(activations_log_dir)
# Check if directory exists
if not log_dir.exists():
return False
# Check for rank files (at least rank_0.pth should exist)
rank_files = list(log_dir.glob("rank_*.pth"))
if not rank_files:
return False
# Check for args.json (created by main process)
args_file = log_dir / "args.json"
has_args_json = args_file.exists()
# Check for completion: if we have rank files and args.json, scoring is complete
if rank_files and has_args_json:
# Add optional completion info for debugging
mprint(f"Found completed scoring in {activations_log_dir}")
mprint(f" - Found {len(rank_files)} rank files")
mprint(f" - Found args.json: {has_args_json}")
return True
return False
def should_skip_scoring_completely(cfg: DictConfig) -> bool:
"""Determine if we should skip scoring entirely (only if 100% complete).
Partial progress should proceed to validate_model for proper resume.
Args:
cfg: Configuration object
Returns:
bool: True if we should skip scoring (100% completed), False if we should run/resume it
"""
# Check if activations_log_dir is specified
if not hasattr(cfg.pruning, "activations_log_dir") or cfg.pruning.activations_log_dir is None:
mprint("No activations_log_dir specified, running scoring")
return False
# Check for force restart flag
force_restart = getattr(cfg.pruning, "force_restart_scoring", False)
if force_restart:
mprint("Force restart flag set, will restart scoring regardless of existing artifacts")
return False
# Get hook configuration to check if resume is mathematically safe
activation_hooks_kwargs = getattr(cfg.pruning, "activation_hooks_kwargs", {})
# Check if scoring is already completed
is_completed = check_scoring_completion(
cfg.pruning.activations_log_dir, activation_hooks_kwargs
)
# Broadcast the result to all processes in distributed mode
if dist.size() > 1:
should_skip = [is_completed] # Use list for mutable object
torch.distributed.broadcast_object_list(should_skip, src=0)
is_completed = should_skip[0]
if is_completed:
mprint("Scoring 100% completed, skipping...")
return is_completed
def launch_score_activations(cfg: DictConfig):
from ..tools.validate_model import validate_model
# Check if we should skip scoring entirely (only if 100% complete)
if should_skip_scoring_completely(cfg):
return
mprint("Starting pruning activation scoring...")
# The checkpoint manager inside validate_model handles all progress tracking
validate_model(args=cfg.pruning)
modelopt/torch/puzzletron/anymodel/README.md
+204
View File
@@ -0,0 +1,204 @@
# AnyModel Guide
This guide explains how to add support for new models in the Puzzletron pipeline.
## Convert model
Convert a HuggingFace model to Puzzletron format.
Step 1: Create Model Descriptor
Extend `ModelDescriptor` and implement `layer_name_predicates()` to define regex patterns for grouping weights into subblocks (embeddings, lm_head, block_N_ffn, block_N_attention).
Key points:
- Find weight names on the model's HuggingFace page → click "Files info" to see the safetensors structure with all tensor names (example: [Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct?show_file_info=model.safetensors.index.json))
See example: [llama_model_descriptor.py](models/llama/llama_model_descriptor.py)
Step 2: Create Converter
Extend `Converter` and implement `create_block_configs_from_main_config()` to create per-layer BlockConfigs from the HuggingFace config.
Key points:
- Import correct HuggingFace config class (e.g., `MistralConfig`, `LlamaConfig`, `Qwen2Config`). Find it in the transformers source: `github.com/huggingface/transformers/tree/main/src/transformers/models/<model_type>/configuration_<model_type>.py`
See example: [llama_converter.py](models/llama/llama_converter.py)
Step 3: Create `models/<model_name>/__init__.py`
Export descriptor and converter classes:
```python
from models.<model_name>.<model_name>_model_descriptor import MyModelDescriptor
from models.<model_name>.<model_name>_converter import MyConverter
```
Step 4: Register in `models/__init__.py`
Add import to trigger factory registration:
```python
from models.<model_name> import *
```
## Usage
```python
from modelopt.torch.puzzletron.anymodel import convert_model
convert_model(
input_dir="path/to/hf_checkpoint",
output_dir="path/to/puzzletron_checkpoint",
converter="model_name",
)
```
## Compress model
Run pruning and compression on a Puzzletron model.
Step 1: Implement ModelDescriptor methods for compression
Add to your `ModelDescriptor`:
- `decoder_layer_cls()` - return the decoder layer class(es) to patch for heterogeneous config support
- `block_config_to_layer_overrides()` - map BlockConfig to layer override dict (see [details](#implementing-block_config_to_layer_overrides))
- `init_rotary_embedding()` - reinitialize rotary embeddings after model loading (see [details](#implementing-init_rotary_embedding))
- `input_embedding_name()` - return the name of the input embedding layer (see [details](#implementing-path-based-methods))
- `output_embedding_name()` - return the name of the output embedding layer (see [details](#implementing-path-based-methods))
- `layer_block_name()` - return the name pattern for decoder layers (see [details](#implementing-path-based-methods))
- `final_norm_name()` - return the name of the final normalization layer (see [details](#implementing-path-based-methods))
- `attn_no_op_post_init()` - replace attention sublayers with no-op modules
- `mlp_no_op_post_init()` - replace MLP sublayers with no-op modules
Step 2: Create FFN Layer Descriptor
Extend `FFNIntermediateLayerDescriptor` to define model-specific paths for FFN pruning hooks (`down_proj_name`, `ffn_prefix_name`, `linear_weight_names`). Derive values from your model's weight names in `layer_name_predicates()`.
See example: [llama_model_descriptor.py](models/llama/llama_model_descriptor.py) → `LlamaFFNIntermediateLayerDescriptor`
Step 3: Configure YAML files
Update the main model config YAML:
- Set `descriptor` to match the name used in `@ModelDescriptorFactory.register_decorator("your_model_name")`
- See example: [llama_3_1_8b_instruct.yaml](../../../../tests/gpu/torch/puzzletron/resources/configs/llama_3_1_8b_instruct/llama_3_1_8b_instruct.yaml)
Update pruning YAML files (`ffn_pruning.yaml`, `expert_pruning.yaml`, etc.):
- Set `pruning_mixin._target_` to the appropriate mixin class
- Set `layer_descriptor._target_` to your layer descriptor class
- Set `hook_class` to the activation hook for scoring
- Set `target_layer` in `activation_hooks_kwargs` to the layer name for hook attachment
- See examples in [configs/llama_3_1_8b_instruct/pruning/](../../../../tests/gpu/torch/puzzletron/resources/configs/llama_3_1_8b_instruct/pruning/)
## End-to-end example
See [test_puzzletron.py](../../../../tests/gpu/torch/puzzletron/test_puzzletron.py) for a complete example that runs both convert and compression steps. For container setup and dependencies needed to run this test, see the [Puzzletron README environment section](../../../../examples/puzzletron/README.md#environment).
---
## Advanced Topics
## Pruning Configuration
### Pruning YAML Structure
Each pruning type has a YAML config with these key fields:
```yaml
pruning_mixin:
_target_: pruning.<type>_pruning_mixin.<MixinClass>
layer_descriptor:
_target_: models.<model>.<descriptor_class>
hook_class: ${get_object:utils.activation_hooks.hooks.<HookClass>}
activation_hooks_kwargs:
method: <method_name>
target_layer: "<layer.name>" # e.g., "mlp.down_proj", "self_attn.o_proj"
```
| Field | Description |
|-------|-------------|
| `pruning_mixin._target_` | Mixin class that orchestrates this pruning type |
| `layer_descriptor._target_` | Model-specific class defining layer paths for hooks |
| `hook_class` | Activation hook class for importance scoring |
| `target_layer` | Layer name (relative to decoder block) where hooks attach |
### Adding a New Hook Class
1. **Implement the hook** under `modelopt/torch/prune/importance_hooks/` (e.g. `base_hooks.py` for generic hooks, `expert_removal_hooks.py` for MoE expert removal):
- Extend an existing hook base class (e.g., `RemoveExpertsIndependentHook` in `expert_removal_hooks.py`)
- Implement required methods (e.g., `get_router_logits_and_routed_experts`)
2. **Register the hook** in the appropriate pruning mixin's `supported_hooks()`:
For FFN pruning (`pruning/ffn_intermediate_pruning_mixin.py`):
```python
def supported_hooks(self) -> List[Type[ActivationsHook]]:
return [IndependentChannelContributionHook, IterativeChannelContributionHook, YourNewHook]
```
For expert removal (`pruning/expert_removal_pruning_mixin.py`):
```python
def supported_hooks(self) -> List[Type[ActivationsHook]]:
return [RankedChoiceVotingHook, ..., YourNewHook]
```
3. **Reference in YAML**:
```yaml
hook_class: ${get_object:utils.activation_hooks.hooks.YourNewHook}
```
### Pruning Types Reference
| Type | Mixin | Example Hooks |
|------|-------|---------------|
| FFN intermediate | [`FFNIntermediatePruningMixIn`](../pruning/ffn_intermediate_pruning_mixin.py) | [`IterativeChannelContributionHook`](../../prune/importance_hooks/base_hooks.py), [`IndependentChannelContributionHook`](../../prune/importance_hooks/base_hooks.py) |
| Expert removal | [`ExpertRemovalPruningMixIn`](../pruning/expert_removal_pruning_mixin.py) | [`NemotronHRemoveExpertsIndependentHook`](../../prune/importance_hooks/expert_removal_hooks.py), [`Qwen3VLRemoveExpertsIndependentHook`](../../prune/importance_hooks/expert_removal_hooks.py) |
| KV heads | [`KVHeadsPruningMixIn`](../pruning/kv_heads_pruning_mixin.py) | [`IndependentKvHeadContributionHook`](../../prune/importance_hooks/base_hooks.py) |
## Implementing `block_config_to_layer_overrides`
Maps Puzzletron's [`BlockConfig`](../block_config.py) fields to HuggingFace config attribute names. Only override attributes that change during pruning:
| BlockConfig Field | HuggingFace Attribute (check `config.json`) |
|-------------------|---------------------------------------------|
| `attention.num_key_value_heads` | `num_key_value_heads` |
| `ffn.intermediate_size` | `intermediate_size` |
| `ffn.moe.num_local_experts` | `num_experts` or `n_routed_experts` (model-specific) |
| `ffn.moe.expert_intermediate_dim` | `moe_intermediate_size` |
**Tip**: Check the model's `config.json` for exact attribute names - they vary between models.
See examples: [qwen3_vl](models/qwen3_vl/qwen3_vl_model_descriptor.py), [nemotron_h](models/nemotron_h/nemotron_h_model_descriptor.py)
---
## Implementing path-based methods
These methods return paths derived from the model's weight names:
- `input_embedding_name()`, `output_embedding_name()`, `layer_block_name()`, `final_norm_name()`
Find them on the model's HuggingFace page → "Files info" → safetensors structure (example: [Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct?show_file_info=model.safetensors.index.json)).
See example: [llama_model_descriptor.py](models/llama/llama_model_descriptor.py)
---
## Implementing `init_rotary_embedding`
Rotary embeddings are computed modules (not saved weights). After model sharding, they need re-initialization on the correct device/dtype.
Look in `github.com/huggingface/transformers/tree/main/src/transformers/models/<model_type>/modeling_<model_type>.py` for:
- `class.*Rotary` — the rotary embedding class name and constructor arguments
- `self.rotary_emb` — the attribute path
See example: [llama_model_descriptor.py](models/llama/llama_model_descriptor.py)
modelopt/torch/puzzletron/anymodel/__init__.py
+43
View File
@@ -0,0 +1,43 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
"""AnyModel: Architecture-agnostic model compression for HuggingFace models.
This module provides a declarative approach to model compression that works with
any HuggingFace model without requiring custom modeling code. Instead of duplicating
HuggingFace modeling classes, AnyModel uses ModelDescriptors that define:
1. Which decoder layer class(es) to patch for heterogeneous configs
2. How to map BlockConfig to layer-specific overrides
3. Weight name patterns for subblock checkpointing
Example usage:
>>> from modelopt.torch.puzzletron.anymodel import convert_model
>>> convert_model(
... input_dir="path/to/hf_checkpoint",
... output_dir="path/to/anymodel_checkpoint",
... converter="llama",
... )
Supported models:
- llama: Llama 2, Llama 3, Llama 3.1, Llama 3.2
- (more to come: qwen2, mistral_small, etc.)
"""
from . import models # trigger factory registration
from .converter import *
from .model_descriptor import *
from .puzzformer import *
modelopt/torch/puzzletron/anymodel/converter/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
"""Converters for transforming HuggingFace models to AnyModel format."""
from .base import *
from .convert_any_model import *
from .converter_factory import *
modelopt/torch/puzzletron/anymodel/converter/base.py
+239
View File
@@ -0,0 +1,239 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
import copy
import fnmatch
import json
import os
import shutil
from abc import ABC, abstractmethod
from collections import defaultdict
from pathlib import Path
from typing import Dict, List
from safetensors.torch import load_file, save_file
from tqdm import tqdm
from transformers import PretrainedConfig
from transformers.integrations.mxfp4 import convert_moe_packed_tensors
from ...block_config import BlockConfig
from ...tools.checkpoint_utils_hf import load_model_config, save_model_config
from ..model_descriptor import ModelDescriptor
__all__ = ["Converter"]
class Converter(ABC):
"""Base class for converting HuggingFace models to Puzzletron/AnyModel format."""
@staticmethod
def _get_weight_map(input_dir: Path) -> Dict[str, str]:
"""Load weight map from checkpoint directory (supports both sharded and single-file models).
Returns a dict mapping parameter names to their safetensors filenames.
"""
index_path = input_dir / "model.safetensors.index.json"
single_file_path = input_dir / "model.safetensors"
if index_path.exists():
# Sharded model
with open(index_path, "r") as f:
index = json.load(f)
return index["weight_map"]
elif single_file_path.exists():
# Single file model - create a synthetic weight map
data = load_file(single_file_path)
return {name: "model.safetensors" for name in data.keys()}
else:
raise FileNotFoundError(
f"Neither {index_path} nor {single_file_path} found. Cannot determine model format."
)
@classmethod
def convert_model_weights(
cls, input_dir: Path, output_dir: Path, descriptor: ModelDescriptor, num_hidden_layers: int
):
"""Convert model weights to subblock format."""
param_to_file = Converter._get_weight_map(input_dir)
all_param_names = list(param_to_file.keys())
# Reverse map: file -> set of params
file_to_params = defaultdict(set)
for name, file in param_to_file.items():
file_to_params[file].add(name)
# Determine subblocks needed
subblocks = descriptor.get_weight_groups(
all_param_names, num_hidden_layers=num_hidden_layers
)
# Output directory
out_dir = output_dir / "subblocks_safetensors"
os.makedirs(out_dir, exist_ok=True)
# New weight index
new_index = {"metadata": {"format": "pt"}, "weight_map": {}}
for subblock, param_names in tqdm(subblocks.items(), desc="Processing subblocks"):
param_files = set(param_to_file[name] for name in param_names)
tensors = {}
# Load only needed files for this subblock
for file in param_files:
data = load_file(os.path.join(input_dir, file))
for name in param_names:
if param_to_file[name] == file and name in data:
converted_name = cls.convert_weight_name(name)
# Convert MoE packed tensors if quantized is mxfp4 //gpt-oss-20b
if getattr(cls, "quantized", None) == "mxfp4":
if name.endswith("_blocks"):
converted_name = converted_name.replace("_blocks", "")
tensors[converted_name] = convert_moe_packed_tensors(
data[name],
data[name.replace("_blocks", "_scales")],
)
elif name.endswith("_scales"):
continue
else:
tensors[converted_name] = data[name]
else:
tensors[converted_name] = data[name]
# Save this subblock
print(f"\n✅ Group: {subblock} ({len(tensors)} layers)")
for layer in tensors.keys():
print(f" - {layer}")
subblock_file = f"{subblock}.safetensors"
save_file(tensors, os.path.join(out_dir, subblock_file))
# Update index
for new_name in tensors.keys():
new_index["weight_map"][new_name] = f"subblocks_safetensors/{subblock_file}"
# Save new index file
with (output_dir / "model.safetensors.index.json").open("w") as f:
json.dump(new_index, f, indent=2)
print(f"✅ Finished saving subblocks and index to {output_dir}")
@classmethod
def convert_configs_in_dirs(
cls,
input_dir: Path,
output_dir: Path,
trust_remote_code: bool = False,
):
"""Convert config and add block_configs."""
config = load_model_config(input_dir, trust_remote_code=trust_remote_code)
block_configs = cls.create_block_configs_from_main_config(config)
out_config = copy.deepcopy(config)
out_config.block_configs = block_configs
save_model_config(out_config, output_dir)
return out_config
@staticmethod
def copy_checkpoint_files(input_dir: Path, output_dir: Path):
"""Copy checkpoint files except model weights (which will be converted)."""
ignore_patterns = [
"model-*.safetensors",
"model.safetensors",
"model.safetensors.index.json",
"subblocks_safetensors",
]
def ignore_func(dir, files):
ignored = set()
for pattern in ignore_patterns:
ignored.update(fnmatch.filter(files, pattern))
return ignored
shutil.copytree(str(input_dir), str(output_dir), ignore=ignore_func, dirs_exist_ok=True)
@classmethod
def convert(
cls,
descriptor: ModelDescriptor,
input_dir: Path,
output_dir: Path,
):
"""Convert a HuggingFace model to AnyModel format.
Args:
descriptor: Model descriptor for the model type.
input_dir: Path to the input HuggingFace checkpoint.
output_dir: Path to the output AnyModel checkpoint.
"""
cls.copy_checkpoint_files(input_dir, output_dir)
trust_remote_code = descriptor.requires_trust_remote_code()
config = cls.convert_configs_in_dirs(
input_dir, output_dir, trust_remote_code=trust_remote_code
)
cls.convert_model_weights(
input_dir, output_dir, descriptor=descriptor, num_hidden_layers=config.num_hidden_layers
)
@staticmethod
@abstractmethod
def create_block_configs_from_main_config(config: PretrainedConfig) -> List[BlockConfig]:
"""Create per-layer BlockConfig list from a HuggingFace model config.
This method extracts layer-specific parameters (e.g., intermediate_size,
num_key_value_heads) from the main model config and creates a BlockConfig
for each layer. These BlockConfigs enable layer-specific pruning and
modifications during the compression pipeline.
Args:
config: HuggingFace PretrainedConfig (e.g., LlamaConfig, Qwen2Config)
Returns:
List of BlockConfig, one per hidden layer. Each BlockConfig contains:
- AttentionConfig: attention settings (no_op, num_key_value_heads)
- FFNConfig: FFN settings (no_op, intermediate_size)
Example:
For a model with uniform layers (e.g., Llama):
return [BlockConfig(...)] * config.num_hidden_layers
For a model with heterogeneous layers (e.g., NemotronH with Mamba/Attention):
return [BlockConfig(...) for layer_idx in range(num_layers)]
"""
raise NotImplementedError
@staticmethod
def convert_weight_name(name: str) -> str:
"""
Convert weight names during checkpoint conversion.
This method can be overridden by subclasses to apply model-specific weight name
transformations when converting checkpoints from HuggingFace format to Puzzletron format.
Default implementation returns the name unchanged (identity function).
Args:
name: Original weight name from HuggingFace checkpoint
Returns:
Converted weight name for Puzzletron format
Example:
For Qwen2.5-VL, this converts:
- visual.* → model.visual.*
- model.* → model.language_model.*
"""
return name
modelopt/torch/puzzletron/anymodel/converter/convert_any_model.py
+68
View File
@@ -0,0 +1,68 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
"""Convert a HuggingFace model to AnyModel format."""
from pathlib import Path
from ..model_descriptor import ModelDescriptorFactory
from .base import Converter
from .converter_factory import ConverterFactory
__all__ = ["convert_model"]
def convert_model(
input_dir: str,
output_dir: str,
converter: Converter | str,
):
"""Convert a HuggingFace model to AnyModel format.
This function converts a HuggingFace checkpoint to the AnyModel format used
for compression. The conversion process:
1. Copies non-weight files (config, tokenizer, etc.)
2. Creates block_configs for each layer
3. Reorganizes weights into subblock checkpoints
Args:
input_dir: Path to the input HuggingFace checkpoint directory.
output_dir: Path to the output AnyModel checkpoint directory.
converter: Either a converter name (e.g., "llama") or a Converter class.
Example:
>>> convert_model(
... input_dir="/path/to/Llama-3.1-8B-Instruct",
... output_dir="/path/to/output/ckpts/teacher",
... converter="llama",
... )
"""
input_dir = Path(input_dir)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# Get descriptor and converter from factories (they use the same name)
descriptor = ModelDescriptorFactory.get(converter)
converter = ConverterFactory.get(converter)
converter.convert(descriptor=descriptor, input_dir=input_dir, output_dir=output_dir)
if __name__ == "__main__":
from fire import Fire
Fire(convert_model)
modelopt/torch/puzzletron/anymodel/converter/converter_factory.py
+75
View File
@@ -0,0 +1,75 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
import inspect
from typing import Callable, Type
from ..model_descriptor import ModelDescriptor
__all__ = ["ConverterFactory"]
class ConverterFactory:
"""Factory for registering and retrieving Converter classes."""
CLASS_MAPPING = {}
@classmethod
def register(cls, **entries: Type):
"""Register converter classes.
Raises:
KeyError: if entry key is already in type_dict and points to a different class.
"""
for cls_name, cls_type in entries.items():
if cls_name in cls.CLASS_MAPPING:
ref = cls.CLASS_MAPPING[cls_name]
# If ref and cls_name point to the same class ignore and don't raise an exception.
if cls_type == ref:
continue
raise KeyError(
f"Could not register `{cls_name}`: {cls_type}, "
f"`{cls_name}` is already registered and points to "
f"`{inspect.getmodule(ref).__name__}.{ref.__name__}`"
)
cls.CLASS_MAPPING[cls_name] = cls_type
@classmethod
def register_decorator(cls, name: str | None) -> Callable:
"""Set up a register decorator.
Args:
name: If specified, the decorated object will be registered with this name.
Returns:
Decorator that registers the callable.
"""
def decorator(cls_type: Type) -> Callable:
"""Register the decorated callable."""
cls_name = name if name is not None else cls_type.__name__
cls.register(**{cls_name: cls_type})
return cls_type
return decorator
@classmethod
def get(cls, value: str | ModelDescriptor):
"""Get a registered converter by name or return the converter if already resolved."""
if isinstance(value, str):
if value in cls.CLASS_MAPPING:
return cls.CLASS_MAPPING[value]
return value
modelopt/torch/puzzletron/anymodel/model_descriptor/__init__.py
+18
View File
@@ -0,0 +1,18 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
"""Model descriptors for defining model-specific properties and layer naming conventions."""
from .base import *
from .model_descriptor_factory import *
modelopt/torch/puzzletron/anymodel/model_descriptor/base.py
+253
View File
@@ -0,0 +1,253 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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 re
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import Any, Dict, Iterable, List, Type
import torch.nn as nn
from ...block_config import BlockConfig
from ...utils.dummy_modules import DummyBlock
__all__ = ["ModelDescriptor"]
class ModelDescriptor(ABC):
@staticmethod
@abstractmethod
def decoder_layer_cls() -> Type[nn.Module] | List[Type[nn.Module]]:
"""Decoder layer class types to patch for heterogeneous config support.
In most cases this class will hold as attributes both FFN & attention layers.
Returns:
nn.Module class type or a list if several class types should be patched.
"""
raise NotImplementedError
@staticmethod
@abstractmethod
def block_config_to_layer_overrides(block_config: BlockConfig) -> Dict[str, Any]:
"""Map between BlockConfig and layer config overrides.
These overrides are consumed by a specific decoder layer and by the whole model.
Usage can be seen in `deci_x_patcher` under the method `_patched_decoder_layer_init`.
Example implementation to override the FFN intermediate size of a block:
>>> def block_config_to_layer_overrides(block_config: BlockConfig) -> Dict[str, Any]:
>>> return {"intermediate_size": block_config.ffn.intermediate_size}
"""
raise NotImplementedError
@staticmethod
def requires_trust_remote_code() -> bool:
"""Whether this model descriptor requires trust_remote_code=True for loading.
Models that use custom code (e.g., via auto_map in config) should override
this to return True.
Returns:
True if trust_remote_code=True is required, False otherwise.
"""
return False
@staticmethod
def mlp_no_op_post_init(decoder_layer: nn.Module):
"""Post-init callback to alter a decoder layer so that FFN/mlp subblock performs as no-op.
It is recommended to use the utils modules from `no_op.py` to replace layers to dummy
counterparts.
Example for replacing a layernorm layer with identity:
>>> decoder_layer.post_attention_layernorm = Same()
Example for replacing an MLP layer with zeroes (zeroes since hidden_states are added to
the residuals hidden_states so a no-op implementation will leave residual the same):
>>> decoder_layer.mlp = MatchingZeros()
In case the MLP layer to replace returns multiple outputs i.e `hidden_states, _ = self.mlp()`,
use the util method `return_tuple_of_size` to return trailing None values:
>>> decoder_layer.mlp = return_tuple_of_size(MatchingZeros, size=2)()
"""
raise NotImplementedError
@staticmethod
def attn_no_op_post_init(decoder_layer: nn.Module):
"""Post-init callback to alter a decoder layer so that Attention subblock performs as no-op.
It is recommended to use the utils modules from `no_op.py` to replace layers to dummy
counterparts.
Example for replacing a layernorm layer with identity:
>>> decoder_layer.post_attention_layernorm = Same()
Example for replacing an attention layer with zeroes:
>>> decoder_layer.self_attn = MatchingZeros()
In case the attention layer returns multiple outputs i.e `hidden_states, _ = self.self_attn()`,
use the util method `return_tuple_of_size` to return trailing None values:
>>> decoder_layer.self_attn = return_tuple_of_size(MatchingZeros, size=2)()
"""
raise NotImplementedError
@staticmethod
@abstractmethod
def init_rotary_embedding(model, runtime):
"""Re-initiate the rotary embeddings based on an existing model.
In puzzletron we initiate a sharded model by first creating a meta model then replacing
to the actual device by loading the state_dict with the real weights.
Rotary embeddings frequencies are tensor buffers that are created dynamically during init
and are not part of the model state_dict, so cannot be restored after a meta device
initialization.
"""
raise NotImplementedError
@staticmethod
@abstractmethod
def input_embedding_name():
"""Return the name of the input embedding layer."""
raise NotImplementedError
@staticmethod
@abstractmethod
def output_embedding_name():
"""Return the name of the output embedding layer."""
raise NotImplementedError
@staticmethod
@abstractmethod
def final_norm_name():
"""Return the name of the final normalization layer."""
raise NotImplementedError
@staticmethod
@abstractmethod
def layer_block_name(index: int):
"""Return the name of the decoder layer at the given index."""
raise NotImplementedError
@staticmethod
@abstractmethod
def layer_name_predicates(num_layers: int) -> Dict[str, re.Pattern]:
"""Return predicates for grouping model weights to support subblock checkpointing.
For every group name return a regex predicate whether a layer name is part of the group.
Returns:
Dictionary of group name to regex pattern predicate.
"""
raise NotImplementedError
@staticmethod
def uses_autocast() -> bool:
"""Whether this model supports torch.autocast.
Some models (e.g., Qwen3-VL MoE) have dtype bugs under autocast.
Override and return False for models that do not support autocast.
"""
return True
@staticmethod
def get_language_model_config(config):
"""Get the language model config from a PretrainedConfig.
For regular LM models, returns the config itself.
For VL/multimodal models with nested configs, override to return the
language model portion (e.g., config.text_config for Qwen-VL).
"""
return config
@staticmethod
def truncate_pattern_for_subblock(
lm_config: Any, parent_layer_index: int | None = None
) -> None:
"""Adjust per-layer config fields so a single-layer model represents the correct layer type.
The default implementation handles ``hybrid_override_pattern`` for
hybrid architectures. It is a no-op when the field is absent.
Override if a model uses a different pattern alphabet.
"""
pattern = getattr(lm_config, "hybrid_override_pattern", None)
if not pattern:
return
# Strip cosmetic pipe separators (e.g. "M|-|*" -> "M-*") before indexing.
pattern = pattern.replace("|", "")
if not pattern:
raise ValueError(
f"hybrid_override_pattern is set but contains no layer-type characters "
f"(original: {lm_config.hybrid_override_pattern!r})"
)
if parent_layer_index is not None and 0 <= parent_layer_index < len(pattern):
lm_config.hybrid_override_pattern = pattern[parent_layer_index]
return
lm_config.hybrid_override_pattern = pattern[0]
@classmethod
def create_dummy_block(cls, original_layer: nn.Module, block_index: int) -> nn.Module:
"""Create a dummy block to replace a layer for sharded model initialization."""
return DummyBlock(block_index=block_index)
@classmethod
def mlp_no_op_supported(cls) -> bool:
"""Check whether `mlp_no_op_post_init` is overridden for mlp no-op support."""
method_name = ModelDescriptor.mlp_no_op_post_init.__name__
return getattr(cls, method_name) is not getattr(ModelDescriptor, method_name)
@classmethod
def attn_no_op_supported(cls):
"""Check whether `attn_no_op_post_init` is overridden for attention no-op support."""
method_name = ModelDescriptor.attn_no_op_post_init.__name__
return getattr(cls, method_name) is not getattr(ModelDescriptor, method_name)
@classmethod
def get_weight_groups(
cls, layer_names: Iterable[str], num_hidden_layers: int
) -> Dict[str, List[str]]:
"""Group model weights to support the puzzle subblock checkpointing format.
This method uses the abstract method `layer_name_predicates` by default.
Args:
layer_names: state_dict layer names of the model.
num_hidden_layers: number of decoder layers in the model.
Returns:
Dictionary of group names to list of layer names per group, e.g.:
>>> {
... "embedding": ["model.embed_tokens.weight"],
... "lm_head": ["lm_head.weight", "model.norm.weight"],
... "block_0_ffn": ["model.layers.0.mlp.down_proj", ...],
... "block_0_attention": ["model.layers.0.self_attn.q_proj", ...],
... }
"""
weight_groups = defaultdict(list)
for name in layer_names:
for group, pattern in cls.layer_name_predicates(num_hidden_layers).items():
if pattern.match(name):
weight_groups[group].append(name)
break
else:
raise ValueError(f"Couldn't find a match for {name}")
return weight_groups
modelopt/torch/puzzletron/anymodel/model_descriptor/model_descriptor_factory.py
+124
View File
@@ -0,0 +1,124 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
import inspect
from typing import Callable, Type
from transformers import AutoConfig
from ...tools.checkpoint_utils_hf import force_cache_dynamic_modules
from .base import ModelDescriptor
__all__ = ["ModelDescriptorFactory", "resolve_descriptor_from_pretrained"]
# Map from HuggingFace config.model_type (in checkpoint config.json) to ModelDescriptorFactory name.
# Local to this script; add entries when supporting new model types for auto-detection.
_MODEL_TYPE_TO_DESCRIPTOR = {
"llama": "llama",
"mistral": "mistral_small",
"qwen2": "qwen2",
"qwen3": "qwen3",
"nemotron_h": "nemotron_h",
"nemotron_h_v2": "nemotron_h_v2",
"gpt_oss_20b": "gpt_oss_20b",
}
def resolve_descriptor_from_pretrained(pretrained: str, trust_remote_code: bool = False):
"""Resolve the model descriptor by loading the checkpoint config and mapping model_type.
Args:
pretrained: Path to a pretrained model checkpoint or HuggingFace model identifier.
trust_remote_code: If True, allows execution of custom code from the model repository.
This is a security risk if the model source is untrusted. Only set to True if you
trust the source of the model. Defaults to False for security.
Returns:
The resolved ModelDescriptor class for the detected model type.
Raises:
ValueError: If pretrained is not provided or if the model type cannot be auto-detected.
"""
config = AutoConfig.from_pretrained(pretrained, trust_remote_code=trust_remote_code)
force_cache_dynamic_modules(config, pretrained, trust_remote_code=trust_remote_code)
model_type = getattr(config, "model_type", None)
if model_type and model_type in _MODEL_TYPE_TO_DESCRIPTOR:
detected = _MODEL_TYPE_TO_DESCRIPTOR[model_type]
print(
f"[resolve_descriptor_from_pretrained] Auto-detected model_type='{model_type}' → descriptor='{detected}'"
)
return ModelDescriptorFactory.get(detected)
known = sorted(_MODEL_TYPE_TO_DESCRIPTOR.keys())
raise ValueError(
f"Cannot auto-detect descriptor for model_type='{model_type}'. "
f"Known model types: {known}. Add this model_type to _MODEL_TYPE_TO_DESCRIPTOR if supported."
)
class ModelDescriptorFactory:
"""Factory for registering and retrieving ModelDescriptor classes."""
CLASS_MAPPING = {}
@classmethod
def register(cls, **entries: Type):
"""Register model descriptor classes.
Raises:
KeyError: if entry key is already in type_dict and points to a different class.
"""
for cls_name, cls_type in entries.items():
if cls_name in cls.CLASS_MAPPING:
ref = cls.CLASS_MAPPING[cls_name]
# If ref and cls_name point to the same class ignore and don't raise an exception.
if cls_type == ref:
continue
raise KeyError(
f"Could not register `{cls_name}`: {cls_type}, "
f"`{cls_name}` is already registered and points to "
f"`{inspect.getmodule(ref).__name__}.{ref.__name__}`"
)
cls.CLASS_MAPPING[cls_name] = cls_type
@classmethod
def register_decorator(cls, name: str | None) -> Callable:
"""Set up a register decorator.
Args:
name: If specified, the decorated object will be registered with this name.
Returns:
Decorator that registers the callable.
"""
def decorator(cls_type: Type) -> Callable:
"""Register the decorated callable."""
cls_name = name if name is not None else cls_type.__name__
cls.register(**{cls_name: cls_type})
return cls_type
return decorator
@classmethod
def get(cls, value: str | ModelDescriptor):
"""Get a registered model descriptor by name or return the descriptor if already resolved."""
if isinstance(value, str):
if value in cls.CLASS_MAPPING:
return cls.CLASS_MAPPING[value]
return value
modelopt/torch/puzzletron/anymodel/models/__init__.py
+29
View File
@@ -0,0 +1,29 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
from packaging.version import Version as _Version
from transformers import __version__ as _transformers_version
# Import models to trigger factory registration
from .gpt_oss import *
from .llama import *
from .mistral_small import *
from .nemotron_h import *
from .nemotron_h_v2 import *
from .qwen2 import *
from .qwen3 import *
if _Version(_transformers_version) >= _Version("4.57.0"):
from .qwen3_vl import *
modelopt/torch/puzzletron/anymodel/models/gpt_oss/__init__.py
+22
View File
@@ -0,0 +1,22 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""GPT-OSS model support for AnyModel."""
from .gpt_oss_converter import *
from .gpt_oss_model_descriptor import *
modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_converter.py
+71
View File
@@ -0,0 +1,71 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
"""GPT-OSS-20B converter for AnyModel compression."""
from typing import List
from transformers import PretrainedConfig
from ....block_config import AttentionConfig, BlockConfig, FFNConfig, MoEConfig
from ...converter import Converter, ConverterFactory
__all__ = ["GptOssConverter"]
@ConverterFactory.register_decorator("gpt_oss")
class GptOssConverter(Converter):
"""Converter for GPT-OSS models to AnyModel format.
GPT-OSS is a pure MoE model with 32/128 experts per layer and 4/16 active experts.
All layers use MoE FFN (no standard dense FFN layers).
"""
quantized = "mxfp4"
@staticmethod
def create_block_configs_from_main_config(config: PretrainedConfig) -> List[BlockConfig]:
"""Create block configs for GPT-OSS layers.
GPT-OSS uses MoE for all FFN layers with:
- 32/128 local experts (num_local_experts)
- 4/16 active experts per token (experts_per_token)
- No dense/standard FFN layers
"""
num_hidden_layers = config.num_hidden_layers
num_local_experts = config.num_local_experts
experts_per_token = config.experts_per_token
intermediate_size = config.intermediate_size
block_configs = []
for layer_idx in range(num_hidden_layers):
block_config = BlockConfig(
attention=AttentionConfig(
no_op=False, num_key_value_heads=config.num_key_value_heads
),
ffn=FFNConfig(
no_op=False,
intermediate_size=None, # MoE doesn't use this field
moe=MoEConfig(
num_local_experts=num_local_experts,
num_experts_per_tok=experts_per_token,
expert_intermediate_dim=intermediate_size,
),
),
).to_dict()
block_configs.append(block_config)
return block_configs
modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_model_descriptor.py
+232
View File
@@ -0,0 +1,232 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
# mypy: ignore-errors
"""GPT-OSS model descriptor for AnyModel compression."""
import re
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Type
import torch.nn as nn
from transformers.models.gpt_oss.modeling_gpt_oss import GptOssDecoderLayer, GptOssRotaryEmbedding
from ....block_config import BlockConfig
from ....pruning.expert_removal_pruning_mixin import (
ExpertRemovalLayerDescriptor,
ExpertRemovalPruningMixIn,
)
# Expert removal is supported for unquantized models (test models).
# Production models use MXFP4 quantized MoE with combined tensors
# (gate_up_proj_blocks, down_proj_blocks), which is not yet supported.
from ....pruning.pruning_mixin import PruningMixIn
from ....utils.dummy_modules import DummyBlock
from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory
from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size
__all__ = ["GptOssModelDescriptor", "GptOssExpertRemovalLayerDescriptor"]
@ModelDescriptorFactory.register_decorator("gpt_oss")
class GptOssModelDescriptor(ModelDescriptor):
"""Model descriptor for GPT-OSS (pure MoE model)."""
_DECODER_LAYER_CLS: Type[nn.Module] = None
@classmethod
def create_dummy_block(cls, original_layer: GptOssDecoderLayer, block_index: int) -> nn.Module:
dummy_block = DummyBlock(block_index=block_index)
# Required by `GptOssModel.forward` in transformers<5.4
if hasattr(original_layer, "attention_type"):
dummy_block.attention_type = original_layer.attention_type
return dummy_block
@staticmethod
def decoder_layer_cls():
"""Get the decoder layer class for GPT-OSS models.
GPT-OSS is a standard transformers model in recent versions.
Import directly from transformers.models.gpt_oss.modeling_gpt_oss.
"""
return GptOssDecoderLayer
@staticmethod
def block_config_to_layer_overrides(block_config: BlockConfig):
"""Map BlockConfig to layer constructor overrides."""
override_kwargs = {}
if block_config.attention.num_key_value_heads is not None:
override_kwargs["num_key_value_heads"] = block_config.attention.num_key_value_heads
if block_config.ffn.moe is not None:
override_kwargs["moe_intermediate_size"] = block_config.ffn.moe.expert_intermediate_dim
override_kwargs["num_local_experts"] = block_config.ffn.moe.num_local_experts
override_kwargs["num_experts_per_tok"] = block_config.ffn.moe.num_experts_per_tok
return override_kwargs
@staticmethod
def attn_no_op_post_init(decoder_layer):
"""Replace attention sublayers with no-op modules."""
decoder_layer.input_layernorm = Same()
decoder_layer.self_attn = return_tuple_of_size(MatchingZeros, size=2)()
@staticmethod
def mlp_no_op_post_init(decoder_layer):
"""Replace MLP sublayers with no-op modules.
Note: GPT-OSS MoE layers return (hidden_states, router_scores), so we need
to return a tuple of 2 values.
"""
decoder_layer.post_attention_layernorm = Same()
decoder_layer.mlp = return_tuple_of_size(MatchingZeros, size=2)()
@staticmethod
def init_rotary_embedding(model, runtime):
"""Initialize rotary embeddings on the correct device."""
# GPT-OSS uses RoPE with YARN scaling
model.model.rotary_emb = GptOssRotaryEmbedding(
config=model.config,
device=runtime.device,
)
@staticmethod
def input_embedding_name():
return "model.embed_tokens"
@staticmethod
def output_embedding_name():
return "lm_head"
@staticmethod
def final_norm_name():
return "model.norm"
@staticmethod
def layer_block_name(index: int):
return f"model.layers.{index}"
@staticmethod
def layer_name_predicates(num_layers: int) -> Dict[str, re.Pattern]:
"""Define regex patterns for grouping weights into subblocks."""
layer_name_patterns = {
"embeddings": re.compile(r"^model\.embed_tokens\.weight$"),
"lm_head": re.compile(r"^(model\.norm\.weight|lm_head\.weight)$"),
}
def build_ffn_predicates() -> Dict[str, re.Pattern]:
"""FFN is MoE in GPT-OSS with MXFP4 quantization."""
return {
f"block_{layer_idx}_ffn": re.compile(
rf"^model\.layers\.{layer_idx}\."
r"(post_attention_layernorm\.weight"
r"|mlp\.router\.weight"
r"|mlp\.router\.bias"
r"|mlp\.experts\.(gate_up_proj|down_proj)(_(bias|blocks|scales))?)$"
)
for layer_idx in range(num_layers)
}
def build_attention_predicates() -> Dict[str, re.Pattern]:
return {
f"block_{layer_idx}_attention": re.compile(
rf"^model\.layers\.{layer_idx}\."
r"(input_layernorm\.weight"
r"|self_attn\.q_proj\.weight"
r"|self_attn\.q_proj\.bias"
r"|self_attn\.k_proj\.weight"
r"|self_attn\.k_proj\.bias"
r"|self_attn\.v_proj\.weight"
r"|self_attn\.v_proj\.bias"
r"|self_attn\.o_proj\.weight"
r"|self_attn\.o_proj\.bias"
r"|self_attn\.sinks)$"
)
for layer_idx in range(num_layers)
}
layer_name_patterns.update(
**build_ffn_predicates(),
**build_attention_predicates(),
)
return layer_name_patterns
@staticmethod
def pruning_mixins() -> Dict[str, PruningMixIn]:
"""Return available pruning mixins for GPT-OSS.
Note: Expert removal works for unquantized models (test models).
Production models use MXFP4 quantization which is not yet supported.
"""
return {"expert_removal": ExpertRemovalPruningMixIn(GptOssExpertRemovalLayerDescriptor())}
@dataclass
class GptOssExpertRemovalLayerDescriptor(ExpertRemovalLayerDescriptor):
"""
GPT-OSS MoE layer descriptor for expert removal.
Note: This only works for unquantized models (e.g., test models).
Production GPT-OSS models use MXFP4 quantization with fused experts
(_blocks, _scales, _bias), which requires a different approach.
Structure:
- Router: mlp.router with .weight and .bias
- Experts: mlp.experts.{idx}.{gate_up_proj,down_proj} with .weight and .bias
"""
target_name: str = "mlp"
moe_prefix_name: str = "model.layers.{layer_idx}.mlp"
expert_prefix_name: str = "experts"
# Router has both weight and bias
router_weights: List[str] = field(default_factory=lambda: ["router.weight"])
router_biases: List[str] = field(default_factory=lambda: ["router.bias"])
# Fused format: experts stored as single tensors
is_fused_experts: bool = True
# Fused format: single tensors containing all experts (test models)
fused_expert_weights: List[str] = field(
default_factory=lambda: [
"experts.gate_up_proj",
"experts.gate_up_proj_bias",
"experts.down_proj",
"experts.down_proj_bias",
]
)
# Not used for fused format, but kept for compatibility
expert_weights: List[str] = field(default_factory=lambda: ["gate_up_proj", "down_proj"])
expert_biases: List[str] = field(
default_factory=lambda: ["gate_up_proj_bias", "down_proj_bias"]
)
def get_modules_names_to_hook(self, model) -> List[Tuple[int, str]]:
target_class_name = "GptOssTopKRouter"
module_names_to_hook = []
for module_name, module in model.named_modules():
if (
module_name.endswith(self.target_name)
and module.__class__.__name__ == target_class_name
):
module_names_to_hook.append(
(self.block_idx_from_module_name(module_name), module_name)
)
return module_names_to_hook
modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py
+526
View File
@@ -0,0 +1,526 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
"""
Create a HuggingFace checkpoint with MXFP4 MoE weights from the original gpt-oss-120b model.
This script:
1. Copies non-MoE weights from the student model (trained attention, embeddings, etc.)
2. Extracts MoE expert weights from the original gpt-oss-120b in MXFP4 format
3. Deduces expert mappings by comparing weights
4. Outputs a new pruned (heterogeneous) checkpoint with PACKED MXFP4 expert weights
"""
import argparse
import json
import os
import shutil
from typing import Any, Dict, List, Tuple
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from tqdm import tqdm
from transformers.integrations.mxfp4 import convert_moe_packed_tensors
__all__ = []
def deduce_experts_for_layer(
layer: int,
original_path: str,
original_index: Dict,
student_path: str,
) -> Tuple[List[int], int, int]:
"""
Deduce which original experts match the student experts by comparing weights.
Compares dequantized MXFP4 weights from the original model against the student
model's BF16 weights using L2 distance. Finds the best 1-to-1 matching.
Args:
layer: Layer index
original_path: Path to original model
original_index: Original model's safetensors index
student_path: Path to student model
num_student_experts: Number of experts in student model (if None, auto-detect)
Returns:
Tuple of (expert_indices, num_student_experts, num_original_experts)
"""
# Load original tensors
orig_tensors = load_layer_tensors(original_path, layer, original_index)
mlp1_blocks = orig_tensors[f"model.layers.{layer}.mlp.experts.gate_up_proj_blocks"]
mlp1_scales = orig_tensors[f"model.layers.{layer}.mlp.experts.gate_up_proj_scales"]
mlp2_blocks = orig_tensors[f"model.layers.{layer}.mlp.experts.down_proj_blocks"]
mlp2_scales = orig_tensors[f"model.layers.{layer}.mlp.experts.down_proj_scales"]
num_original_experts = mlp1_blocks.shape[0]
# Load student tensors
student_subblocks = os.path.join(student_path, "subblocks_safetensors")
student_ffn = os.path.join(student_subblocks, f"block_{layer}_ffn.safetensors")
if not os.path.exists(student_ffn):
print(f"FFN file not found at {student_ffn} - fallback to no_op")
return [], 0, num_original_experts
student_experts = {}
with safe_open(student_ffn, framework="pt") as f:
for key in f.keys():
if "experts" in key or "router" in key:
student_experts[key] = f.get_tensor(key)
# Auto-detect number of student experts
num_student_experts = student_experts[f"model.layers.{layer}.mlp.experts.gate_up_proj"].size(0)
print(
f" Layer {layer}: Comparing {num_student_experts} student experts against {num_original_experts} original experts"
)
# Pre-dequantize all original experts once (optimization)
print(f" Pre-dequantizing {num_original_experts} original experts...")
deqexpert_mlp1 = convert_moe_packed_tensors(mlp1_blocks, mlp1_scales).cpu()
deqexpert_mlp2 = convert_moe_packed_tensors(mlp2_blocks, mlp2_scales).cpu()
original_experts_dequant = []
for orig_idx in range(num_original_experts):
original_experts_dequant.append(
{"up": deqexpert_mlp1[orig_idx], "down": deqexpert_mlp2[orig_idx]}
)
# For each student expert, find best matching original expert
experts_to_keep = []
used_original_indices = set()
# Number of values to use for quick comparison (tune this)
quick_compare_size = 8
# Number of candidates to keep for full comparison
top_k_candidates = min(10, num_original_experts)
for student_idx in range(num_student_experts):
# Get student expert weights
prefix = f"model.layers.{layer}.mlp"
student_up = student_experts.get(f"{prefix}.experts.gate_up_proj")[student_idx] # type: ignore[index]
student_down = student_experts.get(f"{prefix}.experts.down_proj")[student_idx] # type: ignore[index]
# if student_gate is None or student_up is None or student_down is None:
if student_up is None or student_down is None:
raise ValueError(
f"Missing student expert weights for layer {layer} expert {student_idx}"
)
# Step 1: Quick filtering using first N values
candidate_scores = []
for orig_idx in range(num_original_experts):
if orig_idx in used_original_indices:
continue
orig_expert = original_experts_dequant[orig_idx]
up_quick = (
(
orig_expert["up"].flatten()[:quick_compare_size]
- student_up.float().flatten()[:quick_compare_size]
)
.pow(2)
.mean()
.sqrt()
)
down_quick = (
(
orig_expert["down"].flatten()[:quick_compare_size]
- student_down.float().flatten()[:quick_compare_size]
)
.pow(2)
.mean()
.sqrt()
)
quick_score = (up_quick + down_quick) / 2.0
candidate_scores.append((orig_idx, quick_score.item()))
# Step 2: Get top-k candidates based on quick comparison
candidate_scores.sort(key=lambda x: x[1])
top_candidates = [idx for idx, _ in candidate_scores[:top_k_candidates]]
# Step 3: Full comparison only on top candidates
best_match_idx = None
best_match_score = float("inf")
for orig_idx in top_candidates:
orig_expert = original_experts_dequant[orig_idx]
# Full comparison across all values
up_diff = (orig_expert["up"] - student_up.float()).pow(2).mean().sqrt()
down_diff = (orig_expert["down"] - student_down.float()).pow(2).mean().sqrt()
score = (up_diff + down_diff) / 2.0
if score < best_match_score:
best_match_score = score
best_match_idx = orig_idx
if best_match_idx is None:
raise ValueError(
f"Could not find match for student expert {student_idx} in layer {layer}"
)
experts_to_keep.append(best_match_idx)
used_original_indices.add(best_match_idx)
print(
f" Student expert {student_idx} -> Original expert {best_match_idx} (RMSE: {best_match_score:.6f})"
)
return experts_to_keep, num_student_experts, num_original_experts
def load_original_index(path: str) -> Dict[str, Any]:
"""Load the original model's safetensors index."""
with open(path, "r") as f:
return json.load(f)
def load_layer_tensors(original_path: str, layer: int, index: Dict) -> Dict[str, torch.Tensor]:
"""Load all MoE-related tensors for a layer, potentially from multiple files."""
keys_to_load = [
f"model.layers.{layer}.mlp.experts.gate_up_proj_blocks",
f"model.layers.{layer}.mlp.experts.gate_up_proj_scales",
f"model.layers.{layer}.mlp.experts.gate_up_proj_bias",
f"model.layers.{layer}.mlp.experts.down_proj_blocks",
f"model.layers.{layer}.mlp.experts.down_proj_scales",
f"model.layers.{layer}.mlp.experts.down_proj_bias",
f"model.layers.{layer}.mlp.router.weight", # Router weight
f"model.layers.{layer}.mlp.router.bias", # Router bias
]
# Group by file
file_to_keys = {}
for key in keys_to_load:
if key in index["weight_map"]:
filename = index["weight_map"][key]
if filename not in file_to_keys:
file_to_keys[filename] = []
file_to_keys[filename].append(key)
# Load from each file
tensors = {}
for filename, keys in file_to_keys.items():
filepath = os.path.join(original_path, filename)
with safe_open(filepath, framework="pt") as f:
for key in keys:
tensors[key] = f.get_tensor(key)
return tensors
def copy_non_moe_weights(student_path: str, output_path: str, num_layers: int) -> Dict[str, str]:
"""
Copy non-MoE weights from student model.
Returns weight_map for the new index.
"""
weight_map = {}
subblocks_dir = os.path.join(output_path, "subblocks_safetensors")
os.makedirs(subblocks_dir, exist_ok=True)
student_subblocks = os.path.join(student_path, "subblocks_safetensors")
# Copy embeddings
src_emb = os.path.join(student_subblocks, "embeddings.safetensors")
dst_emb = os.path.join(subblocks_dir, "embeddings.safetensors")
shutil.copy2(src_emb, dst_emb)
with safe_open(src_emb, framework="pt") as f:
for key in f.keys():
weight_map[key] = "subblocks_safetensors/embeddings.safetensors"
# Copy lm_head
src_head = os.path.join(student_subblocks, "lm_head.safetensors")
dst_head = os.path.join(subblocks_dir, "lm_head.safetensors")
shutil.copy2(src_head, dst_head)
with safe_open(src_head, framework="pt") as f:
for key in f.keys():
weight_map[key] = "subblocks_safetensors/lm_head.safetensors"
# Copy attention blocks
for layer in range(num_layers):
src_attn = os.path.join(student_subblocks, f"block_{layer}_attention.safetensors")
dst_attn = os.path.join(subblocks_dir, f"block_{layer}_attention.safetensors")
shutil.copy2(src_attn, dst_attn)
with safe_open(src_attn, framework="pt") as f:
for key in f.keys():
weight_map[key] = f"subblocks_safetensors/block_{layer}_attention.safetensors"
return weight_map
def process_single_layer(
layer: int,
original_path: str,
original_index: Dict,
student_path: str,
output_path: str,
experts_to_keep: List[int],
) -> Tuple[Dict[str, str], List[str]]:
"""
Process a single layer - loads tensors from potentially multiple files.
Returns (weight_map, verification_errors).
"""
weight_map = {}
verification_errors = []
subblocks_dir = os.path.join(output_path, "subblocks_safetensors")
student_subblocks = os.path.join(student_path, "subblocks_safetensors")
# Load all tensors for this layer (may come from multiple files)
orig_tensors = load_layer_tensors(original_path, layer, original_index)
# Load student FFN file
student_ffn = os.path.join(student_subblocks, f"block_{layer}_ffn.safetensors")
tensors_to_save = {}
student_tensors = {}
with safe_open(student_ffn, framework="pt") as f:
for key in f.keys():
tensor = f.get_tensor(key)
if "experts" not in key and "router" not in key:
# Copy norm weights
tensors_to_save[key] = tensor
# Get router from original model, sliced to kept experts
orig_router_weight = orig_tensors[f"model.layers.{layer}.mlp.router.weight"]
orig_router_bias = orig_tensors[f"model.layers.{layer}.mlp.router.bias"]
kept_indices_tensor = torch.tensor(experts_to_keep, dtype=torch.long)
sliced_router_weight = orig_router_weight[kept_indices_tensor]
sliced_router_bias = orig_router_bias[kept_indices_tensor]
tensors_to_save[f"model.layers.{layer}.mlp.router.weight"] = sliced_router_weight
tensors_to_save[f"model.layers.{layer}.mlp.router.bias"] = sliced_router_bias
# Get MoE tensors
mlp1_blocks = orig_tensors[f"model.layers.{layer}.mlp.experts.gate_up_proj_blocks"]
mlp1_scales = orig_tensors[f"model.layers.{layer}.mlp.experts.gate_up_proj_scales"]
mlp2_blocks = orig_tensors[f"model.layers.{layer}.mlp.experts.down_proj_blocks"]
mlp2_scales = orig_tensors[f"model.layers.{layer}.mlp.experts.down_proj_scales"]
mlp1_bias = orig_tensors[f"model.layers.{layer}.mlp.experts.gate_up_proj_bias"]
mlp2_bias = orig_tensors[f"model.layers.{layer}.mlp.experts.down_proj_bias"]
tensors_to_save[f"model.layers.{layer}.mlp.experts.gate_up_proj_blocks"] = mlp1_blocks[
kept_indices_tensor
]
tensors_to_save[f"model.layers.{layer}.mlp.experts.gate_up_proj_scales"] = mlp1_scales[
kept_indices_tensor
]
tensors_to_save[f"model.layers.{layer}.mlp.experts.gate_up_proj_bias"] = mlp1_bias[
kept_indices_tensor
]
tensors_to_save[f"model.layers.{layer}.mlp.experts.down_proj_blocks"] = mlp2_blocks[
kept_indices_tensor
]
tensors_to_save[f"model.layers.{layer}.mlp.experts.down_proj_scales"] = mlp2_scales[
kept_indices_tensor
]
tensors_to_save[f"model.layers.{layer}.mlp.experts.down_proj_bias"] = mlp2_bias[
kept_indices_tensor
]
# Save the FFN file
output_file = os.path.join(subblocks_dir, f"block_{layer}_ffn.safetensors")
save_file(tensors_to_save, output_file)
# Build weight map
for key in tensors_to_save.keys():
weight_map[key] = f"subblocks_safetensors/block_{layer}_ffn.safetensors"
return weight_map, verification_errors
def copy_config_files(student_path: str, output_path: str):
"""Copy configuration files from student model and update config.json."""
files_to_copy = [
"tokenizer.json",
"tokenizer_config.json",
"special_tokens_map.json",
"chat_template.jinja",
]
# Also copy transformers compatibility files
if os.path.exists(student_path):
for f in os.listdir(student_path):
if f.startswith("transformers_"):
files_to_copy.append(f)
for filename in files_to_copy:
src = os.path.join(student_path, filename)
dst = os.path.join(output_path, filename)
# Try student path first
if os.path.exists(src):
try:
shutil.copy2(src, dst)
continue
except PermissionError:
pass
# If we get here, file doesn't exist or permission denied
if not os.path.exists(dst):
print(f" Warning: Could not copy {filename}")
# Update config.json for DeciGptOssForCausalLM with MXFP4
src_config = os.path.join(student_path, "config.json")
if not os.path.exists(src_config):
raise FileNotFoundError(f"config.json not found at {src_config}")
with open(src_config, "r") as f:
config = json.load(f) # type: ignore[arg-type]
# Set architecture to DeciGptOssForCausalLM for MXFP4 support
config["architectures"] = ["DeciGptOssForCausalLM"]
# Add quantization_config so vllm calls _load_weights_mxfp4
config["quantization_config"] = {
"quant_method": "mxfp4",
"modules_to_not_convert": [
"model.layers.*.self_attn",
"model.layers.*.mlp.router",
"model.embed_tokens",
"lm_head",
],
}
dst_config = os.path.join(output_path, "config.json")
with open(dst_config, "w") as f:
json.dump(config, f, indent=2) # type: ignore[arg-type]
def main():
parser = argparse.ArgumentParser(description="Create MXFP4 checkpoint from student model")
parser.add_argument(
"--student-path", type=str, required=True, help="Path to student model checkpoint"
)
parser.add_argument(
"--original-path",
type=str,
required=True,
help="Path to original gpt-oss-120b model with MXFP4 weights",
)
parser.add_argument(
"--output-path", type=str, required=True, help="Output path for the new checkpoint"
)
parser.add_argument("--num-layers", type=int, default=36, help="Number of transformer layers")
args = parser.parse_args()
print(f"Creating MXFP4 checkpoint...")
print(f" Student model: {args.student_path}")
print(f" Original model: {args.original_path}")
print(f" Output: {args.output_path}")
# Load original model index
original_index = load_original_index(
os.path.join(args.original_path, "model.safetensors.index.json")
)
print("\nDeducing expert mappings by comparing weights...")
experts_to_keep = []
layer_statistics = [] # Store (num_student, num_original) for each layer
for layer in range(args.num_layers):
layer_experts, num_student, num_original = deduce_experts_for_layer(
layer,
args.original_path,
original_index,
args.student_path,
)
experts_to_keep.append(layer_experts)
layer_statistics.append((num_student, num_original))
# Print statistics
print(f"\n{'=' * 70}")
print("EXPERT DEDUCTION STATISTICS")
print(f"{'=' * 70}")
print(f"{'Layer':<8} {'Student Experts':<18} {'Original Experts':<18} {'Kept %':<10}")
print(f"{'-' * 70}")
total_student = 0
total_original = 0
for layer, (num_student, num_original) in enumerate(layer_statistics):
percentage = (num_student / num_original * 100) if num_original > 0 else 0
print(f"{layer:<8} {num_student:<18} {num_original:<18} {percentage:<10.2f}")
total_student += num_student
total_original += num_original
print(f"{'-' * 70}")
avg_percentage = (total_student / total_original * 100) if total_original > 0 else 0
print(f"{'TOTAL':<8} {total_student:<18} {total_original:<18} {avg_percentage:<10.2f}")
print(f"{'=' * 70}")
print(f"\n Deduced experts_to_keep mapping for {len(experts_to_keep)} layers")
# Create output directory
os.makedirs(args.output_path, exist_ok=True)
os.makedirs(os.path.join(args.output_path, "subblocks_safetensors"), exist_ok=True)
# Copy config files
print("Copying configuration files...")
copy_config_files(args.student_path, args.output_path)
# Save experts_to_keep.json
experts_to_keep_output = os.path.join(args.output_path, "experts_to_keep.json")
with open(experts_to_keep_output, "w") as f:
json.dump(experts_to_keep, f, indent=2)
print(f" Saved experts_to_keep mapping to {experts_to_keep_output}")
# Copy non-MoE weights (embeddings, attention, lm_head)
print("Copying non-MoE weights...")
weight_map = copy_non_moe_weights(args.student_path, args.output_path, args.num_layers)
# Load weights per layer (handles multi-file loading)
print(f"Processing {args.num_layers} layers...")
all_verification_errors = []
# Process each layer
for layer in tqdm(range(args.num_layers), desc="Processing layers"):
if len(experts_to_keep[layer]) == 0:
print(f"Layer {layer} has no experts to keep - ffn->no_op")
continue
layer_weight_map, layer_errors = process_single_layer(
layer,
args.original_path,
original_index,
args.student_path,
args.output_path,
experts_to_keep[layer],
)
weight_map.update(layer_weight_map)
all_verification_errors.extend(layer_errors)
# Calculate total size
total_size = 0
subblocks_dir = os.path.join(args.output_path, "subblocks_safetensors")
for filename in os.listdir(subblocks_dir):
filepath = os.path.join(subblocks_dir, filename)
total_size += os.path.getsize(filepath)
# Create model.safetensors.index.json
index = {"metadata": {"total_size": total_size}, "weight_map": weight_map}
index_path = os.path.join(args.output_path, "model.safetensors.index.json")
with open(index_path, "w") as f:
json.dump(index, f, indent=2)
print(f"\nCheckpoint created successfully at: {args.output_path}")
print(f"Total size: {total_size / 1e9:.2f} GB")
if __name__ == "__main__":
main()
modelopt/torch/puzzletron/anymodel/models/llama/__init__.py
+17
View File
@@ -0,0 +1,17 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
from .llama_converter import *
from .llama_model_descriptor import *

Some files were not shown because too many files have changed in this diff Show More