From 2f8c3b776ab213043bbf5b71b13c6e5b34a00eb1 Mon Sep 17 00:00:00 2001 From: arham766 Date: Mon, 6 Jul 2026 02:51:46 -0700 Subject: [PATCH] test: targeted unit tests for codecov-uncovered lines across eight modules Consolidates the standalone test-suite PRs into a single PR per maintainer feedback in #1902: every test here exercises lines the codecov report marks uncovered, asserts meaningful logic rather than coverage for its own sake, and parametrization clusters are deduplicated. Modules targeted (codecov coverage before): - torch/export/postprocess.py (26.6%): TP/PP split layouts, padding math, quantization-scale resharding - torch/utils/graph.py (26.4%): subgraph matching, memoization, arity/connectivity mismatches - torch/quantization/calib/bias.py (70.0%): method dispatch, running averages, bias round-trips, reset - torch/utils/logging.py (74.3%): tqdm silencing, tty ANSI wrapping, capture_io, warning filters - torch/utils/distributed.py (81.4%): FileLock wait loop, master_only, rank fallbacks - torch/distill/losses.py (81.9%): MGD forward math, MFT threshold branch - torch/quantization/conversion.py (87.0%): SVDQuant entrypoint, restore error paths, SequentialQuantizer edge branches, context manager module-type restoration - torch/utils/robust_json.py (87.0%): encoder special-type fallbacks Also folds in the test_mode_registry assertion fix (was #1921). Lines that remain uncovered need multi-rank process groups or are unreachable defensive branches; they are enumerated in the PR body. Signed-off-by: arham766 --- tests/unit/torch/distill/test_losses.py | 117 ++++ tests/unit/torch/export/test_postprocess.py | 556 ++++++++++++++++++ tests/unit/torch/opt/test_mode_registry.py | 20 +- .../torch/quantization/test_bias_calib.py | 106 ++++ .../torch/quantization/test_conversion.py | 246 ++++++++ tests/unit/torch/utils/test_distributed.py | 86 +++ tests/unit/torch/utils/test_graph.py | 180 ++++++ tests/unit/torch/utils/test_logging.py | 91 +++ tests/unit/torch/utils/test_robust_json.py | 51 ++ 9 files changed, 1451 insertions(+), 2 deletions(-) create mode 100644 tests/unit/torch/distill/test_losses.py create mode 100644 tests/unit/torch/export/test_postprocess.py create mode 100644 tests/unit/torch/quantization/test_bias_calib.py create mode 100644 tests/unit/torch/quantization/test_conversion.py create mode 100644 tests/unit/torch/utils/test_distributed.py create mode 100644 tests/unit/torch/utils/test_graph.py create mode 100644 tests/unit/torch/utils/test_logging.py create mode 100644 tests/unit/torch/utils/test_robust_json.py diff --git a/tests/unit/torch/distill/test_losses.py b/tests/unit/torch/distill/test_losses.py new file mode 100644 index 00000000000..9ae2313c255 --- /dev/null +++ b/tests/unit/torch/distill/test_losses.py @@ -0,0 +1,117 @@ +# 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. + +"""Targeted unit tests for distillation losses not exercised elsewhere. + +tests/unit/torch/distill/test_distill.py wires losses through ``mtd.convert`` +but never runs MGDLoss's forward path nor MFTLoss's ``apply_threshold_to_all=False`` +branch; these tests cover that logic directly with hand-computed expectations. +""" + +import pytest +import torch +import torch.nn as nn + +from modelopt.torch.distill.losses import MFTLoss, MGDLoss + + +@pytest.fixture(autouse=True) +def deterministic_seed(): + torch.manual_seed(1234) + + +class TestMFTLoss: + def test_threshold_not_applied_to_correct_argmax_rows(self): + # apply_threshold_to_all=False leaves correct-argmax rows untouched while + # still correcting incorrect-argmax rows. + # Row 0: p = [0.7, 0.2, 0.1], label 0 (argmax correct) -> unchanged. + # Row 1: p = [0.6, 0.3, 0.1], label 1 (argmax 0 is wrong): + # mixin = (p_argmax - p_label + thr) / (1 + p_argmax - p_label) + # = (0.6 - 0.3 + 0.2) / (1 + 0.3) = 5/13 + # adjusted = p * (8/13) + one_hot(1) * (5/13) = [4.8, 7.4, 0.8] / 13 + mft = MFTLoss(threshold=0.2) + p = torch.tensor([[0.7, 0.2, 0.1], [0.6, 0.3, 0.1]]) + corrected = mft._prepare_corrected_distributions( + p.log(), torch.tensor([0, 1]), 0.2, apply_threshold_to_all=False + ) + expected = torch.tensor([[0.7, 0.2, 0.1], [4.8 / 13.0, 7.4 / 13.0, 0.8 / 13.0]]) + assert torch.allclose(corrected, expected, atol=1e-5) + + +class TestMGDLoss: + def test_zeroed_generation_gives_mse_against_zero(self): + # Matching channel counts must make align an identity (no 1x1 conv). + # With lambda_mgd=0 the random mask is all ones (rand() > 1 never holds), + # and with all generation weights/biases zeroed the reconstructed feature + # map is exactly zero, so the loss reduces to mse(0, out_t) = mean(out_t^2). + mgd = MGDLoss(4, 4, lambda_mgd=0.0) + assert isinstance(mgd.align, nn.Identity) + for param in mgd.generation.parameters(): + param.data.zero_() + out_s = torch.randn(2, 4, 5, 5) + out_t = torch.randn(2, 4, 5, 5) + loss = mgd(out_s, out_t) + assert torch.isclose(loss, (out_t**2).mean(), atol=1e-6) + + def test_alpha_mgd_scales_loss_exactly(self): + # lambda_mgd=0 makes the forward deterministic (mask is all ones), and + # sharing weights isolates the alpha factor: loss(alpha=3) == 3 * loss(alpha=1). + mgd_a = MGDLoss(4, 4, alpha_mgd=1.0, lambda_mgd=0.0) + mgd_b = MGDLoss(4, 4, alpha_mgd=3.0, lambda_mgd=0.0) + mgd_b.load_state_dict(mgd_a.state_dict()) + out_s = torch.randn(2, 4, 6, 6) + out_t = torch.randn(2, 4, 6, 6) + assert torch.isclose(mgd_b(out_s, out_t), 3.0 * mgd_a(out_s, out_t), atol=1e-6) + + def test_full_mask_ignores_student_features(self): + # lambda_mgd=1.0 zeroes the entire mask (rand() > 0 holds a.s.), so the + # student features are fully masked out and the loss only depends on + # generation(0) vs the teacher features. + mgd = MGDLoss(4, 4, lambda_mgd=1.0) + out_t = torch.randn(1, 4, 3, 3) + loss_a = mgd(torch.randn(1, 4, 3, 3), out_t) + loss_b = mgd(torch.full((1, 4, 3, 3), 7.0), out_t) + assert torch.isclose(loss_a, loss_b, atol=1e-7) + + def test_channel_alignment_and_gradient_flow(self): + # Mismatched channels route the student features through the 1x1 align + # conv, and gradients must reach the student input and every loss-module + # parameter so DistillationModel can train them. + mgd = MGDLoss(3, 5, lambda_mgd=0.0) + out_s = torch.randn(1, 3, 4, 4, requires_grad=True) + out_t = torch.randn(1, 5, 4, 4) + loss = mgd(out_s, out_t) + assert loss.ndim == 0 + assert loss.item() >= 0.0 # MSE is non-negative + loss.backward() + assert out_s.grad is not None + assert mgd.align.weight.grad is not None + assert all(p.grad is not None for p in mgd.generation.parameters()) + + def test_spatial_mismatch_raises(self): + mgd = MGDLoss(4, 4) + with pytest.raises(AssertionError): + mgd(torch.randn(1, 4, 4, 4), torch.randn(1, 4, 5, 5)) + + def test_teacher_features_not_detached(self): + # NOTE: documents current behavior; arguably a bug because MGDLoss does + # not detach the teacher feature map (unlike LogitsDistillationLoss and + # MFTLoss, which detach their targets), so gradients propagate into the + # teacher whenever its features require grad. + mgd = MGDLoss(4, 4, lambda_mgd=0.0) + out_s = torch.randn(1, 4, 3, 3, requires_grad=True) + out_t = torch.randn(1, 4, 3, 3, requires_grad=True) + mgd(out_s, out_t).backward() + assert out_t.grad is not None diff --git a/tests/unit/torch/export/test_postprocess.py b/tests/unit/torch/export/test_postprocess.py new file mode 100644 index 00000000000..b18c60f0e70 --- /dev/null +++ b/tests/unit/torch/export/test_postprocess.py @@ -0,0 +1,556 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the single-process export postprocess logic (TP/PP config splitting, padding, +lm_head quantization update, weight shape validation and tensor postprocessing). + +The multi-rank merge paths (``_merge_model_configs_to_first_tp`` and +``_model_model_configs_to_first_pp``) require multiple distributed processes exchanging +tensors through SharedMemory/NFS and are not covered here. +""" + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.model_config import ( + LINEAR_COLUMN, + LINEAR_GROUP, + LINEAR_ROW, + AttentionConfig, + ConvConfig, + DecoderLayerConfig, + EmbeddingConfig, + ExpertConfig, + LayernormConfig, + LinearConfig, + MLPConfig, + ModelConfig, + RelativeAttentionTableConfig, +) +from modelopt.torch.export.postprocess import ( + _same_tensor, + _shallow_copy_with_field_instantiation, + _split_model_config_for_pp, + _split_model_config_for_tp, + check_weight_shape_valid, + pad_embedding_lm_head, + postprocess_model_config, + postprocess_tensors, + update_lm_head_quantization, + view_as_float8_e4m3fn_if_needed, + view_as_uint8_if_needed, +) + + +def _column_linear(out_features=8, in_features=4, bias=False, **kwargs): + return LinearConfig( + linear_type=LINEAR_COLUMN, + weight=torch.randn(out_features, in_features), + bias=torch.randn(out_features) if bias else None, + **kwargs, + ) + + +def _row_linear(out_features=4, in_features=8, bias=False, **kwargs): + return LinearConfig( + linear_type=LINEAR_ROW, + weight=torch.randn(out_features, in_features), + bias=torch.randn(out_features) if bias else None, + **kwargs, + ) + + +def _make_model_config(num_layers=2, hidden=4, vocab=8): + """Builds a minimal but structurally realistic ModelConfig.""" + torch.manual_seed(0) + layers = [] + for _ in range(num_layers): + attention = AttentionConfig( + qkv=_column_linear(3 * hidden, hidden), + dense=_row_linear(hidden, hidden), + ) + mlp = MLPConfig( + fc=_column_linear(2 * hidden, hidden), + proj=_row_linear(hidden, 2 * hidden), + hidden_act="silu", + ) + layers.append( + DecoderLayerConfig( + attention=attention, + mlp=mlp, + input_layernorm=LayernormConfig(weight=torch.ones(hidden)), + post_layernorm=LayernormConfig(weight=torch.ones(hidden)), + ) + ) + return ModelConfig( + vocab_size=vocab, + vocab_embedding=EmbeddingConfig(weight=torch.randn(vocab, hidden)), + position_embedding=EmbeddingConfig(weight=torch.randn(vocab, hidden)), + ln_embed=LayernormConfig(weight=torch.ones(hidden)), + layers=layers, + ln_f=LayernormConfig(weight=torch.ones(hidden)), + lm_head=_column_linear(vocab, hidden), + ) + + +INT4_BLOCK32_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 32}}, + "enable": True, + }, + { + "quantizer_name": "*input_quantizer", + "cfg": {"num_bits": 8, "axis": None}, + "enable": True, + }, + ], + "algorithm": "max", +} + + +def _make_quantized_lm_head(out_features=100, in_features=64, quant_cfg=INT4_BLOCK32_CFG): + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(in_features, out_features, bias=False)) + mtq.quantize(model, quant_cfg, lambda m: m(torch.randn(2, 8, in_features))) + return model[0] + + +def test_same_tensor(): + assert _same_tensor([]) + assert _same_tensor([None, None, None]) + assert _same_tensor([torch.ones(2, 2), torch.ones(2, 2)]) + assert not _same_tensor([torch.ones(2), torch.zeros(2)]) + + +class TestFp8Views: + def test_fp8_uint8_roundtrip_preserves_bits(self): + t = torch.randn(4, 4).to(torch.float8_e4m3fn) + viewed = view_as_uint8_if_needed(t) + assert viewed.dtype == torch.uint8 + assert viewed.shape == t.shape + roundtrip = view_as_float8_e4m3fn_if_needed(viewed) + assert roundtrip.dtype == torch.float8_e4m3fn + assert torch.equal(roundtrip.view(torch.uint8), t.view(torch.uint8)) + + def test_other_dtypes_pass_through(self): + t = torch.randn(4, 4) + assert view_as_uint8_if_needed(t) is t + assert view_as_float8_e4m3fn_if_needed(t) is t + + +def test_shallow_copy_with_field_instantiation(): + expert = ExpertConfig(fc=_column_linear(), proj=_row_linear()) + clone = _shallow_copy_with_field_instantiation(expert) + # Dataclass fields are new instances, but tensors within them are still shared. + assert clone is not expert + assert clone.fc is not expert.fc + assert clone.proj is not expert.proj + assert clone.fc.weight is expert.fc.weight + # Mutating a cloned field does not leak back to the original. + clone.fc.weight = torch.zeros(1, 1) + assert expert.fc.weight.shape == (8, 4) + + +class TestSplitModelConfigForTp: + def test_column_linear_splits_weight_and_bias_on_dim0(self): + config = _column_linear(out_features=8, in_features=4, bias=True) + weight, bias = config.weight, config.bias + configs = _split_model_config_for_tp(config, 2) + assert len(configs) == 2 + for i, cfg in enumerate(configs): + assert cfg.weight.shape == (4, 4) + assert torch.equal(cfg.weight, weight.chunk(2, dim=0)[i]) + assert torch.equal(cfg.bias, bias.chunk(2, dim=0)[i]) + + def test_column_linear_pads_odd_rows_before_split(self): + config = _column_linear(out_features=5, in_features=4) + configs = _split_model_config_for_tp(config, 2) + assert configs[0].weight.shape == (3, 4) + assert configs[1].weight.shape == (3, 4) + # The last row of the second shard is the zero padding. + assert torch.equal(configs[1].weight[-1], torch.zeros(4)) + # NOTE: documents current behavior; arguably a bug because the input + # (merged) config's weight is reassigned by _split_model_config_for_tp. + assert config.weight.shape == (6, 4) + + def test_row_linear_splits_weight_on_dim1_and_keeps_bias(self): + config = _row_linear(out_features=4, in_features=8, bias=True) + weight, bias = config.weight, config.bias + configs = _split_model_config_for_tp(config, 2) + for i, cfg in enumerate(configs): + assert cfg.weight.shape == (4, 4) + assert torch.equal(cfg.weight, weight.chunk(2, dim=1)[i]) + # Row linear bias is not split; all shards share it. + assert cfg.bias is bias + + def test_tp_disabled_linear_is_not_split(self): + config = _column_linear(out_features=8, in_features=4, tp=False) + configs = _split_model_config_for_tp(config, 2) + assert len(configs) == 2 + for cfg in configs: + assert cfg.weight is config.weight + + def test_column_scaling_factor_split_on_dim0_and_scalar_shared(self): + # Per-channel (fp8) scaling factors split on dim 0, preserving the dtype. + wsf = torch.rand(8, 2).to(torch.float8_e4m3fn) + config = _column_linear(out_features=8, weights_scaling_factor=wsf) + configs = _split_model_config_for_tp(config, 2) + for i, cfg in enumerate(configs): + assert cfg.weights_scaling_factor.dtype == torch.float8_e4m3fn + assert cfg.weights_scaling_factor.shape == (4, 2) + assert torch.equal( + cfg.weights_scaling_factor.view(torch.uint8), + wsf.view(torch.uint8).chunk(2, dim=0)[i], + ) + + # Scalar scaling factors are shared, not split. + scalar = _column_linear(weights_scaling_factor=torch.tensor(0.5)) + for cfg in _split_model_config_for_tp(scalar, 2): + assert cfg.weights_scaling_factor is scalar.weights_scaling_factor + + def test_row_awq_scaling_factors_split(self): + # AWQ row linear: weight scaling factor splits on dim 1 (preserving fp8 dtype), + # prequant scaling factor splits on dim 0. + wsf = torch.rand(8, 4).to(torch.float8_e4m3fn) + psf = torch.rand(16) + config = _row_linear( + out_features=8, + in_features=16, + weights_scaling_factor=wsf, + prequant_scaling_factor=psf, + awq_block_size=4, + ) + configs = _split_model_config_for_tp(config, 2) + for i, cfg in enumerate(configs): + assert cfg.weights_scaling_factor.dtype == torch.float8_e4m3fn + assert torch.equal( + cfg.weights_scaling_factor.view(torch.uint8), + wsf.view(torch.uint8).chunk(2, dim=1)[i], + ) + assert torch.equal(cfg.prequant_scaling_factor, psf.chunk(2, dim=0)[i]) + + def test_row_int8_sq_scaling_factor_not_split(self): + # awq_block_size == 0 means INT8 SQ: per-channel scaling factor stays intact. + wsf = torch.rand(8) + config = _row_linear(out_features=8, in_features=16, weights_scaling_factor=wsf) + configs = _split_model_config_for_tp(config, 2) + for cfg in configs: + assert cfg.weights_scaling_factor is wsf + + def test_group_linear_raises(self): + config = LinearConfig(linear_type=LINEAR_GROUP, weight=torch.randn(8, 4)) + with pytest.raises(AssertionError, match="group linear"): + _split_model_config_for_tp(config, 2) + + def test_conv_config_raises(self): + with pytest.raises(NotImplementedError, match="ConvConfig"): + _split_model_config_for_tp(ConvConfig(weight=torch.randn(4, 4, 3, 3)), 2) + + def test_embedding_split_on_dim0_with_padding(self): + config = EmbeddingConfig(weight=torch.randn(5, 4)) + padded = torch.cat([config.weight, torch.zeros(1, 4)], dim=0) + configs = _split_model_config_for_tp(config, 2) + for i, cfg in enumerate(configs): + assert torch.equal(cfg.weight, padded.chunk(2, dim=0)[i]) + # The last row of the second shard is the zero padding. + assert torch.equal(configs[1].weight[-1], torch.zeros(4)) + + def test_relative_attention_table_split_on_dim0(self): + weight = torch.randn(6, 4) + config = RelativeAttentionTableConfig(weight=weight) + configs = _split_model_config_for_tp(config, 3) + for i, cfg in enumerate(configs): + assert torch.equal(cfg.weight, weight.chunk(3, dim=0)[i]) + + def test_expert_config_split(self): + hidden, moe_hidden, num_experts = 4, 8, 2 + config = ExpertConfig( + fc=LinearConfig( + linear_type=LINEAR_COLUMN, + weight=torch.randn(num_experts, 2 * moe_hidden, hidden), + ), + proj=LinearConfig( + linear_type=LINEAR_ROW, + weight=torch.randn(num_experts, hidden, moe_hidden), + ), + ) + fc_weight, proj_weight = config.fc.weight, config.proj.weight + configs = _split_model_config_for_tp(config, 2) + + # proj (row) is split on dim 2. + for i, cfg in enumerate(configs): + assert torch.equal(cfg.proj.weight, proj_weight.chunk(2, dim=2)[i]) + + # fc holds concatenated [w3; w1]: each is split separately then re-concatenated + # so each shard gets matching halves of w3 and w1. + merged_w3, merged_w1 = torch.chunk(fc_weight, 2, dim=1) + for i, cfg in enumerate(configs): + expected = torch.cat( + [merged_w3.chunk(2, dim=1)[i], merged_w1.chunk(2, dim=1)[i]], dim=1 + ) + assert torch.equal(cfg.fc.weight, expected) + + def test_expert_config_awq_scaling_factors_split(self): + hidden, moe_hidden, num_experts, block = 8, 8, 2, 4 + fc_wsf = torch.rand(num_experts, 2 * moe_hidden, hidden // block) + proj_wsf = torch.rand(num_experts, hidden, moe_hidden // block) + proj_psf = torch.rand(num_experts, moe_hidden) + config = ExpertConfig( + fc=LinearConfig( + linear_type=LINEAR_COLUMN, + weight=torch.randn(num_experts, 2 * moe_hidden, hidden), + weights_scaling_factor=fc_wsf, + awq_block_size=block, + ), + proj=LinearConfig( + linear_type=LINEAR_ROW, + weight=torch.randn(num_experts, hidden, moe_hidden), + weights_scaling_factor=proj_wsf, + prequant_scaling_factor=proj_psf, + awq_block_size=block, + ), + ) + configs = _split_model_config_for_tp(config, 2) + + merged_wsf_w3, merged_wsf_w1 = torch.chunk(fc_wsf, 2, dim=1) + for i, cfg in enumerate(configs): + expected_fc_wsf = torch.cat( + [merged_wsf_w3.chunk(2, dim=1)[i], merged_wsf_w1.chunk(2, dim=1)[i]], dim=1 + ) + assert torch.equal(cfg.fc.weights_scaling_factor, expected_fc_wsf) + assert torch.equal(cfg.proj.weights_scaling_factor, proj_wsf.chunk(2, dim=2)[i]) + assert torch.equal(cfg.proj.prequant_scaling_factor, proj_psf.chunk(2, dim=-1)[i]) + + def test_recursion_through_dataclass_and_list(self): + mlp = MLPConfig( + fc=_column_linear(8, 4), + proj=_row_linear(4, 8), + hidden_act="gelu", + ) + configs = _split_model_config_for_tp([mlp], 2) + assert len(configs) == 2 + for i, cfg in enumerate(configs): + assert isinstance(cfg, list) + assert cfg[0].fc.weight.shape == (4, 4) + assert cfg[0].proj.weight.shape == (4, 4) + assert torch.equal(cfg[0].fc.weight, mlp.fc.weight.chunk(2, dim=0)[i]) + assert cfg[0].hidden_act == "gelu" + + +class TestSplitModelConfigForPp: + def test_layers_partitioned_and_heads_assigned(self): + model_config = _make_model_config(num_layers=4) + all_layers = list(model_config.layers) + configs = _split_model_config_for_pp(model_config, 2) + assert len(configs) == 2 + + first, last = configs + assert first.layers == all_layers[:2] + assert last.layers == all_layers[2:] + + # Only the first PP rank keeps the embeddings. + assert first.vocab_embedding is not None + assert first.position_embedding is not None + assert first.ln_embed is not None + assert last.vocab_embedding is None + assert last.position_embedding is None + assert last.ln_embed is None + + # Only the last PP rank keeps the final norm and lm_head. + assert first.ln_f is None + assert first.lm_head is None + assert last.ln_f is not None + assert last.lm_head is not None + + def test_non_divisible_layers_raise(self): + model_config = _make_model_config(num_layers=3) + with pytest.raises(AssertionError): + _split_model_config_for_pp(model_config, 2) + + +class TestPostprocessModelConfig: + """Single-process cases: training TP == PP == 1, so only the split paths.""" + + def test_split_to_tp2(self): + model_config = _make_model_config(hidden=4, vocab=8) + qkv_weight = model_config.layers[0].attention.qkv.weight + dense_weight = model_config.layers[0].attention.dense.weight + embedding_weight = model_config.vocab_embedding.weight + + configs = postprocess_model_config(model_config, inference_tensor_parallel=2) + assert len(configs) == 2 + assert [cfg.rank for cfg in configs] == [0, 1] + assert all(cfg.tensor_parallel == 2 for cfg in configs) + + for i, cfg in enumerate(configs): + # Column linear split on dim 0. + assert torch.equal(cfg.layers[0].attention.qkv.weight, qkv_weight.chunk(2, dim=0)[i]) + # Row linear split on dim 1. + assert torch.equal( + cfg.layers[0].attention.dense.weight, dense_weight.chunk(2, dim=1)[i] + ) + # Embedding split on dim 0. + assert torch.equal(cfg.vocab_embedding.weight, embedding_weight.chunk(2, dim=0)[i]) + + # Layernorm weights are replicated, not split. + assert torch.equal(configs[0].ln_f.weight, configs[1].ln_f.weight) + assert configs[0].ln_f.weight.shape == (4,) + + def test_split_to_tp2_pp2_rank_assignment(self): + model_config = _make_model_config(num_layers=4) + configs = postprocess_model_config( + model_config, inference_tensor_parallel=2, inference_pipeline_parallel=2 + ) + assert len(configs) == 4 + # Ordered as [tp0/pp0, tp0/pp1, tp1/pp0, tp1/pp1] with rank = tp + pp * tp_size. + assert [cfg.rank for cfg in configs] == [0, 2, 1, 3] + assert all(cfg.tensor_parallel == 2 for cfg in configs) + assert all(cfg.pipeline_parallel == 2 for cfg in configs) + assert all(len(cfg.layers) == 2 for cfg in configs) + + # PP rank 0 keeps the embedding, PP rank 1 keeps the lm_head. + pp0_configs = [configs[0], configs[2]] + pp1_configs = [configs[1], configs[3]] + assert all(cfg.vocab_embedding is not None for cfg in pp0_configs) + assert all(cfg.lm_head is None for cfg in pp0_configs) + assert all(cfg.vocab_embedding is None for cfg in pp1_configs) + assert all(cfg.lm_head is not None for cfg in pp1_configs) + + +class TestPadEmbeddingLmHead: + def test_noop_when_vocab_is_multiple_of_padding_factor(self): + model_config = _make_model_config(vocab=128) + embedding_weight = model_config.vocab_embedding.weight + lm_head_weight = model_config.lm_head.weight + pad_embedding_lm_head(model_config) + assert model_config.vocab_size == 128 + assert model_config.vocab_embedding.weight is embedding_weight + assert model_config.lm_head.weight is lm_head_weight + + def test_pads_vocab_embedding_and_lm_head(self): + model_config = _make_model_config(vocab=100, hidden=4) + model_config.lm_head.bias = torch.randn(100) + original_embedding = model_config.vocab_embedding.weight.clone() + original_lm_head = model_config.lm_head.weight.clone() + original_bias = model_config.lm_head.bias.clone() + + pad_embedding_lm_head(model_config) + + assert model_config.vocab_size == 128 + assert model_config.vocab_embedding.weight.shape == (128, 4) + assert model_config.lm_head.weight.shape == (128, 4) + assert model_config.lm_head.bias.shape == (128,) + + assert torch.equal(model_config.vocab_embedding.weight[:100], original_embedding) + assert torch.equal(model_config.vocab_embedding.weight[100:], torch.zeros(28, 4)) + assert torch.equal(model_config.lm_head.weight[:100], original_lm_head) + assert torch.equal(model_config.lm_head.weight[100:], torch.zeros(28, 4)) + assert torch.equal(model_config.lm_head.bias[:100], original_bias) + assert torch.equal(model_config.lm_head.bias[100:], torch.zeros(28)) + + def test_pads_weights_scaling_factor_with_int4_maxbound(self): + model_config = _make_model_config(vocab=100, hidden=4) + original_wsf = torch.rand(100, 2) + model_config.lm_head.weights_scaling_factor = original_wsf.clone() + + pad_embedding_lm_head(model_config) + + wsf = model_config.lm_head.weights_scaling_factor + assert wsf.shape == (128, 2) + assert torch.equal(wsf[:100], original_wsf) + # Padded rows use 1 / 7.0 (int4 maxbound). + assert torch.allclose(wsf[100:], torch.full((28, 2), 1.0 / 7.0)) + + +class TestUpdateLmHeadQuantization: + def test_plain_linear_is_ignored(self): + # No weight_quantizer/input_quantizer attributes: early return, no error. + update_lm_head_quantization(ModelConfig(), nn.Linear(4, 4)) + + def test_non_divisible_vocab_disables_quantization(self): + # 100 % (32 * 1) != 0 -> quantization must be disabled. + lm_head = _make_quantized_lm_head(out_features=100) + assert lm_head.weight_quantizer.is_enabled + update_lm_head_quantization(ModelConfig(), lm_head, inference_tensor_parallel=1) + assert not lm_head.weight_quantizer.is_enabled + assert not lm_head.input_quantizer.is_enabled + + def test_divisible_vocab_keeps_quantization_and_warns(self): + # 128 % (32 * 1) == 0 -> quantization stays enabled with a warning. + lm_head = _make_quantized_lm_head(out_features=128) + with pytest.warns(UserWarning, match="lm_head quantization"): + update_lm_head_quantization(ModelConfig(), lm_head, inference_tensor_parallel=1) + assert lm_head.weight_quantizer.is_enabled + assert lm_head.input_quantizer.is_enabled + + def test_inference_tp_affects_divisibility(self): + # 96 % (32 * 1) == 0 but 96 % (32 * 2) != 0. + lm_head = _make_quantized_lm_head(out_features=96) + with pytest.warns(UserWarning, match="lm_head quantization"): + update_lm_head_quantization(ModelConfig(), lm_head, inference_tensor_parallel=1) + assert lm_head.weight_quantizer.is_enabled + + update_lm_head_quantization(ModelConfig(), lm_head, inference_tensor_parallel=2) + assert not lm_head.weight_quantizer.is_enabled + + def test_awq_lite_pre_quant_scale_is_removed_on_disable(self): + quant_cfg = { + "quant_cfg": INT4_BLOCK32_CFG["quant_cfg"], + "algorithm": "awq_lite", + } + lm_head = _make_quantized_lm_head(out_features=100, quant_cfg=quant_cfg) + assert hasattr(lm_head.input_quantizer, "_pre_quant_scale") + + update_lm_head_quantization(ModelConfig(), lm_head, inference_tensor_parallel=1) + assert not lm_head.weight_quantizer.is_enabled + assert not hasattr(lm_head.input_quantizer, "_pre_quant_scale") + + +class TestCheckWeightShapeValid: + def test_tp_disabled_config_ignores_inference_tp(self): + # k=6 is not divisible by inference TP 4, but tp=False forces TP 1. + config = _column_linear(out_features=8, in_features=6, tp=False) + check_weight_shape_valid(config, inference_tensor_parallel=4) + + def test_expert_config_valid_and_invalid(self): + valid = ExpertConfig( + fc=LinearConfig(linear_type=LINEAR_COLUMN, weight=torch.randn(2, 8, 16)), + proj=LinearConfig(linear_type=LINEAR_ROW, weight=torch.randn(2, 16, 8)), + ) + check_weight_shape_valid(valid, inference_tensor_parallel=4) + + invalid = ExpertConfig( + fc=LinearConfig( + linear_type=LINEAR_COLUMN, weight=torch.randn(2, 8, 16), awq_block_size=8 + ), + proj=LinearConfig(linear_type=LINEAR_ROW, weight=torch.randn(2, 16, 8)), + ) + with pytest.raises(NotImplementedError, match="block size"): + check_weight_shape_valid(invalid, inference_tensor_parallel=4) + + +def test_postprocess_tensors_clones_views(): + base = torch.randn(4, 4) + view = base[:2] + assert view._is_view() + weights = {"a.weight": view} + postprocess_tensors(weights, torch.float32) + assert not weights["a.weight"]._is_view() + assert torch.equal(weights["a.weight"], base[:2]) diff --git a/tests/unit/torch/opt/test_mode_registry.py b/tests/unit/torch/opt/test_mode_registry.py index 271d5c1d7dc..6c5a3c897cc 100644 --- a/tests/unit/torch/opt/test_mode_registry.py +++ b/tests/unit/torch/opt/test_mode_registry.py @@ -13,6 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + +import pytest + from modelopt.torch.opt.mode import ModeDescriptor, _ModeRegistryCls @@ -47,8 +51,20 @@ def restore(self): registry.remove_mode("a_test_mode") + # after removal, the mode must no longer resolve via any registry assert not _ModeRegistryCls.contained_in_any("a_test_mode") + with pytest.raises(KeyError, match="a_test_mode"): + _ModeRegistryCls.get_from_any("a_test_mode") - del registry + # NOTE: `del registry` alone does NOT remove the registry from + # `_ModeRegistryCls._all_registries`: the class-level list holds a strong reference, so the + # local `del` never drops the refcount to zero and `__del__` (which would remove it from the + # list) is unreachable. Remove it explicitly so this test leaves no residue behind for other + # tests sharing the process. Since `__del__` unconditionally calls `list.remove`, it would + # then raise ValueError once the orphaned object is garbage collected, so neutralize it for + # the actual destruction. + _ModeRegistryCls._all_registries.remove(registry) + with mock.patch.object(_ModeRegistryCls, "__del__", return_value=None): + del registry # refcount now drops to zero -> patched `__del__` runs here - assert "test_registry" not in _ModeRegistryCls._all_registries + assert all(r._registry_name != "test_registry" for r in _ModeRegistryCls._all_registries) diff --git a/tests/unit/torch/quantization/test_bias_calib.py b/tests/unit/torch/quantization/test_bias_calib.py new file mode 100644 index 00000000000..7c3b5e7ed60 --- /dev/null +++ b/tests/unit/torch/quantization/test_bias_calib.py @@ -0,0 +1,106 @@ +# 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. + +"""Targeted unit tests for modelopt.torch.quantization.calib.bias. + +Only exercises the paths not covered elsewhere: per-tensor bias math and method +dispatch, bias subtraction/addition, the running-average aggregation of the mean +method, dynamic bias computation, and calibrator reset. +""" + +import torch + +from modelopt.torch.quantization.calib.bias import ( + BiasCalibrator, + add_bias, + compute_bias, + subtract_bias, +) + + +class TestBiasFunctions: + def test_compute_bias_dispatch_per_tensor(self): + # [0, 1, 8]: mean = 3, (max + min) / 2 = 4 -- distinguishes the two methods. + x = torch.tensor([0.0, 1.0, 8.0]) + mean_bias = compute_bias(x, None, method="mean") + maxmin_bias = compute_bias(x, None, method="max_min") + assert mean_bias.shape == () + assert maxmin_bias.shape == () + assert mean_bias.item() == 3.0 + assert maxmin_bias.item() == 4.0 + + def test_subtract_add_bias_round_trip(self): + x = torch.tensor([[1.0, 3.0], [5.0, 9.0]]) + bias = torch.tensor([[2.0], [7.0]]) + centered = subtract_bias(x, bias) + assert torch.equal(centered, torch.tensor([[-1.0, 1.0], [-2.0, 2.0]])) + restored = add_bias(centered, bias) + assert restored.shape == x.shape + assert torch.equal(restored, x) + + +class TestBiasCalibrator: + def test_mean_running_average_weights_collects_equally(self): + # Aggregation is a running average over *collect calls*, not elements: + # collect([1, 2, 3]) -> bias = 2 + # collect([4]) -> bias = (2 * 1 + 4) / 2 = 3 + # collect([9]) -> bias = (3 * 2 + 9) / 3 = 5 + # An element-weighted mean of all 5 values would be 3.8 instead. + calibrator = BiasCalibrator(method="mean", axis=None) + assert calibrator.compute_bias() is None # nothing collected yet + calibrator.collect(torch.tensor([1.0, 2.0, 3.0])) + assert calibrator.compute_bias().item() == 2.0 + calibrator.collect(torch.tensor([4.0])) + assert calibrator.compute_bias().item() == 3.0 + calibrator.collect(torch.tensor([9.0])) + assert calibrator.compute_bias().item() == 5.0 + assert calibrator._cnt == 3 + + def test_mean_running_average_preserves_dtype(self): + # The running average is computed in float32 for stability but cast back. + calibrator = BiasCalibrator(method="mean", axis=None) + calibrator.collect(torch.tensor([1.0, 3.0], dtype=torch.float16)) # mean 2 + calibrator.collect(torch.tensor([5.0], dtype=torch.float16)) # (2 + 5) / 2 = 3.5 + bias = calibrator.compute_bias() + assert bias.dtype == torch.float16 + assert bias.item() == 3.5 + + def test_compute_dynamic_bias_is_stateless_for_both_methods(self): + # [0, 1, 8]: mean = 3, (max + min) / 2 = 4 + x = torch.tensor([0.0, 1.0, 8.0]) + mean_calib = BiasCalibrator(method="mean", axis=None) + mean_calib.collect(torch.tensor([100.0, 200.0])) + assert mean_calib.compute_dynamic_bias(x).item() == 3.0 # collected history ignored + assert mean_calib.compute_bias().item() == 150.0 # collected state untouched + maxmin_calib = BiasCalibrator(method="max_min", axis=None) + assert maxmin_calib.compute_dynamic_bias(x).item() == 4.0 + + def test_reset_clears_state_for_both_methods(self): + mean_calib = BiasCalibrator(method="mean", axis=None) + mean_calib.collect(torch.tensor([1.0, 2.0, 3.0])) + mean_calib.collect(torch.tensor([10.0])) + maxmin_calib = BiasCalibrator(method="max_min", axis=None) + maxmin_calib.collect(torch.tensor([-100.0, 100.0])) + for calibrator in (mean_calib, maxmin_calib): + calibrator.reset() + assert calibrator.compute_bias() is None + assert calibrator._calib_max is None + assert calibrator._calib_min is None + assert calibrator._cnt == 0 + # No history left: fresh collects match a brand-new calibrator. + mean_calib.collect(torch.tensor([10.0, 20.0])) + assert mean_calib.compute_bias().item() == 15.0 + maxmin_calib.collect(torch.tensor([0.0, 2.0])) + assert maxmin_calib.compute_bias().item() == 1.0 # old extrema forgotten diff --git a/tests/unit/torch/quantization/test_conversion.py b/tests/unit/torch/quantization/test_conversion.py new file mode 100644 index 00000000000..faa3e211026 --- /dev/null +++ b/tests/unit/torch/quantization/test_conversion.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Targeted unit tests for modelopt.torch.quantization.conversion. + +Only exercises code paths not covered elsewhere in the unit suite: the SVDQuant +conversion entrypoint, quantizer-state restore error paths, parent-class lookup +failure, SequentialQuantizer edge branches, the temporary-config context +manager's module-type restoration, and the deprecated set_quantizer_attribute shim. +""" + +import pytest +import torch.nn as nn +from _test_utils.torch.quantization.models import SimpleLinear + +from modelopt.torch.opt.conversion import ApplyModeError, ModeloptStateManager +from modelopt.torch.quantization.config import QuantizeConfig, QuantizerAttributeConfig +from modelopt.torch.quantization.conversion import ( + convert_to_quantized_model_svdquant, + quantizer_state, + replace_quant_module, + restore_quantizer_state, + set_quantizer_attribute, + set_quantizer_attributes_full, + set_quantizer_attributes_partial, + set_quantizer_by_cfg, + set_quantizer_by_cfg_context, + update_quantize_metadata, +) +from modelopt.torch.quantization.nn import SequentialQuantizer, SVDQuantLinear, TensorQuantizer + +# SimpleLinear has three nn.Linear layers -> three weight quantizers after replacement. +SIMPLE_LINEAR_NUM_WEIGHT_QUANTIZERS = 3 + + +def _replaced(): + """Return a fresh SimpleLinear with quantized modules inserted (no calibration).""" + model = SimpleLinear() + replace_quant_module(model) + return model + + +def _weight_quantizers(model): + """TensorQuantizer weight quantizers, guaranteed non-empty (guards vacuous loops).""" + found = { + n: m + for n, m in model.named_modules() + if isinstance(m, TensorQuantizer) and n.endswith("weight_quantizer") + } + assert len(found) == SIMPLE_LINEAR_NUM_WEIGHT_QUANTIZERS + return found + + +def _seq_weight_quantizers(model): + """SequentialQuantizer containers on the weight path (subclass nn.Sequential, not + TensorQuantizer, so _weight_quantizers cannot see them).""" + return { + n: m + for n, m in model.named_modules() + if isinstance(m, SequentialQuantizer) and n.endswith("weight_quantizer") + } + + +def _make_saved_metadata(): + """A configured model plus the metadata dict a checkpoint would carry.""" + model = _replaced() + set_quantizer_by_cfg( + model, + [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}}, + ], + ) + metadata = {} + update_quantize_metadata(model, QuantizeConfig(), metadata) + return model, metadata + + +class TestSvdquantConversion: + def test_convert_replaces_quant_linears_and_records_metadata(self): + model = SimpleLinear() + ModeloptStateManager(model, init_state=True) + replace_quant_module(model) # SVDQuant conversion runs on an already-quantized model + config = QuantizeConfig( + quant_cfg=[{"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}}] + ) + converted, metadata = convert_to_quantized_model_svdquant(model, config) + assert converted is model # conversion is in-place + linears = [m for m in model.modules() if isinstance(m, nn.Linear)] + assert len(linears) == SIMPLE_LINEAR_NUM_WEIGHT_QUANTIZERS + assert all(isinstance(m, SVDQuantLinear) for m in linears) + # The quant_cfg is applied to the quantizers of the SVDQuant modules... + for module in _weight_quantizers(model).values(): + assert module.num_bits == 4 + assert module.axis == 0 + # ...and the metadata records the resulting quantizer state. + assert metadata["quantizer_state"] == quantizer_state(model) + + +class TestRestoreQuantizerStateErrors: + # NOTE: the happy-path restore round-trip is already covered by existing unit tests; + # only the key-mismatch error branches are pinned here. + + def test_unmatched_checkpoint_keys_raise(self): + """Keys present in the checkpoint but missing from the model are 'unmatched'.""" + _, metadata = _make_saved_metadata() + metadata["quantizer_state"]["bogus.weight_quantizer"] = {} + with pytest.raises(ApplyModeError, match=r"Unmatched keys.*bogus\.weight_quantizer"): + restore_quantizer_state(_replaced(), QuantizeConfig(), metadata) + + def test_extra_model_keys_raise(self): + """Quantizers in the model that have no checkpoint entry are 'extra'.""" + _, metadata = _make_saved_metadata() + del metadata["quantizer_state"]["net.0.weight_quantizer"] + with pytest.raises(ApplyModeError, match=r"Extra keys.*net\.0\.weight_quantizer"): + restore_quantizer_state(_replaced(), QuantizeConfig(), metadata) + + +def test_unknown_parent_class_raises(): + model = _replaced() + with pytest.raises(ValueError, match="not found in QuantModuleRegistry"): + set_quantizer_by_cfg( + model, + [{"parent_class": "nn.NoSuchModule", "quantizer_name": "*", "enable": False}], + ) + + +class TestSequentialQuantizerEdgeBranches: + # Callable filters are used below so that only the SequentialQuantizer containers match; + # a "*weight_quantizer" wildcard would also reach the containers' ".0/.1" children via + # the fused-experts name normalization and hit different branches. + + def test_full_list_length_mismatch_warns_and_assigns_partially(self): + model = _replaced() + set_quantizer_attributes_full( + model, + lambda name: name.endswith("weight_quantizer"), + [QuantizerAttributeConfig(num_bits=4), QuantizerAttributeConfig(num_bits=8)], + ) + with pytest.warns(UserWarning, match="does not match the number"): + set_quantizer_attributes_full( + model, + lambda name: name.endswith("weight_quantizer"), + [QuantizerAttributeConfig(num_bits=2)], + ) + containers = _seq_weight_quantizers(model) + assert len(containers) == SIMPLE_LINEAR_NUM_WEIGHT_QUANTIZERS + for module in containers.values(): + # Partial assignment: only the first sub-quantizer got the new config. + assert [q.num_bits for q in module] == [2, 8] + + def test_partial_list_on_plain_tensor_quantizer_raises(self): + model = _replaced() + with pytest.raises(ValueError, match="not a SequentialQuantizer"): + set_quantizer_attributes_partial( + model, "*input_quantizer", [{"enable": False}, {"enable": True}] + ) + + def test_partial_list_applies_per_position_on_sequential(self): + model = _replaced() + set_quantizer_attributes_full( + model, + "*weight_quantizer", + [QuantizerAttributeConfig(num_bits=4), QuantizerAttributeConfig(num_bits=8)], + ) + set_quantizer_attributes_partial( + model, + lambda name: name.endswith("weight_quantizer"), + [{"enable": False}, {"enable": True}], + ) + containers = _seq_weight_quantizers(model) + assert len(containers) == SIMPLE_LINEAR_NUM_WEIGHT_QUANTIZERS + for module in containers.values(): + assert not module[0].is_enabled + assert module[1].is_enabled + assert [q.num_bits for q in module] == [4, 8] # merged, not replaced + + +class TestSetQuantizerByCfgContextTypeRestore: + # NOTE: plain attribute save/restore of the context manager is covered elsewhere; + # these tests pin the module-type re-creation branches of the exit handler. + + def test_sequential_downgrade_is_reverted_on_exit(self): + """A single-cfg entry downgrades a SequentialQuantizer; exit re-creates it.""" + model = _replaced() + set_quantizer_by_cfg( + model, + [ + { + "quantizer_name": "*weight_quantizer", + "cfg": [{"num_bits": 4}, {"num_bits": 8, "axis": 0}], + } + ], + ) + with set_quantizer_by_cfg_context( + model, [{"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8}}] + ): + wq = model.net[0].weight_quantizer + assert isinstance(wq, TensorQuantizer) + assert not isinstance(wq, SequentialQuantizer) + assert wq.num_bits == 8 + wq = model.net[0].weight_quantizer + assert isinstance(wq, SequentialQuantizer) + assert len(wq) == 2 + assert [q.num_bits for q in wq] == [4, 8] + assert wq[1].axis == 0 + + def test_manual_upgrade_inside_context_is_reverted(self): + """Even a manual TensorQuantizer->SequentialQuantizer swap in the body is undone.""" + model = _replaced() + with set_quantizer_by_cfg_context(model, []): + set_quantizer_attributes_full( + model, + "*weight_quantizer", + [QuantizerAttributeConfig(num_bits=4), QuantizerAttributeConfig(num_bits=8)], + ) + assert isinstance(model.net[0].weight_quantizer, SequentialQuantizer) + wq = model.net[0].weight_quantizer + assert isinstance(wq, TensorQuantizer) + assert not isinstance(wq, SequentialQuantizer) + assert wq.num_bits == 8 + assert wq.axis == 0 # original default weight attributes restored + + +def test_deprecated_set_quantizer_attribute_warns_and_merges(): + model = _replaced() + set_quantizer_attributes_full( + model, "*weight_quantizer", QuantizerAttributeConfig(num_bits=8, axis=0) + ) + with pytest.warns(DeprecationWarning, match="set_quantizer_attribute is deprecated"): + set_quantizer_attribute(model, "*weight_quantizer", {"num_bits": 4}) + for module in _weight_quantizers(model).values(): + assert module.num_bits == 4 + assert module.axis == 0 # merged like set_quantizer_attributes_partial diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py new file mode 100644 index 00000000000..8edb56d023a --- /dev/null +++ b/tests/unit/torch/utils/test_distributed.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +# Single-process contracts only: everything here must hold without an initialized +# torch.distributed process group. + +import os +import threading + +import pytest +from torch import nn + +import modelopt.torch.utils.distributed as dist + + +def test_local_rank_fallback_warns_and_uses_global_rank(monkeypatch): + monkeypatch.delenv("LOCAL_RANK", raising=False) + with pytest.warns(UserWarning, match="LOCAL_RANK"): + assert dist.local_rank() == 0 + + +def test_master_only_runs_and_returns_result(): + @dist.master_only + def compute(): + return {"x": 1} + + assert compute() == {"x": 1} + + +@pytest.mark.parametrize("group", [None, -1]) +def test_process_group_rank_is_minus_one_when_uninitialized(group): + assert dist.DistributedProcessGroup(group).rank() == -1 + + +def test_get_group_returns_none_when_uninitialized(): + assert dist.get_group([0]) is None + + +def test_is_dtensor_sharded_false_for_plain_model(): + assert not dist.is_dtensor_sharded(nn.Linear(2, 2)) + + +def test_filelock_context_creates_and_removes_lockfile(tmp_path): + lock_path = str(tmp_path / "my.lock") + with dist.FileLock(lock_path): + assert os.path.exists(lock_path) + assert not os.path.exists(lock_path) + + +def test_filelock_try_acquire_conflict(tmp_path): + lock_path = str(tmp_path / "my.lock") + first = dist.FileLock(lock_path) + second = dist.FileLock(lock_path) + assert first.try_acquire() + assert not second.try_acquire() # already held + first.release() + assert second.try_acquire() + second.release() + assert not os.path.exists(lock_path) + + +def test_filelock_all_acquire_waits_for_release(tmp_path): + lock_path = str(tmp_path / "my.lock") + first = dist.FileLock(lock_path) + assert first.try_acquire() + releaser = threading.Timer(0.2, first.release) + releaser.start() + try: + # blocks in wait() polling until the timer releases the first lock + with dist.FileLock(lock_path, all_acquire=True, poll_time=0.02): + assert os.path.exists(lock_path) + assert not os.path.exists(lock_path) + finally: + releaser.join() diff --git a/tests/unit/torch/utils/test_graph.py b/tests/unit/torch/utils/test_graph.py new file mode 100644 index 00000000000..1af14aaedbd --- /dev/null +++ b/tests/unit/torch/utils/test_graph.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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 pytest +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.utils.graph import match + + +class LinearRelu(nn.Module): + def __init__(self, features=8): + super().__init__() + self.fc = nn.Linear(features, features) + self.act = nn.ReLU() + + def forward(self, x): + return self.act(self.fc(x)) + + +class LinearSigmoid(nn.Module): + def __init__(self, features=8): + super().__init__() + self.fc = nn.Linear(features, features) + self.act = nn.Sigmoid() + + def forward(self, x): + return self.act(self.fc(x)) + + +class LinearFuncRelu(nn.Module): + """Same math as LinearRelu but relu as call_function instead of call_module.""" + + def __init__(self, features=8): + super().__init__() + self.fc = nn.Linear(features, features) + + def forward(self, x): + return F.relu(self.fc(x)) + + +class LinearMethodRelu(nn.Module): + """relu invoked as a tensor method (call_method node).""" + + def __init__(self, features=8): + super().__init__() + self.fc = nn.Linear(features, features) + + def forward(self, x): + return self.fc(x).relu() + + +class JustLinear(nn.Module): + def __init__(self, features=8): + super().__init__() + self.fc = nn.Linear(features, features) + + def forward(self, x): + return self.fc(x) + + +class DiamondShared(nn.Module): + """The fc output is consumed by two branches (a node is revisited during matching).""" + + def __init__(self): + super().__init__() + self.fc = nn.Linear(4, 4) + self.relu = nn.ReLU() + self.sig = nn.Sigmoid() + + def forward(self, x): + y = self.fc(x) + return self.relu(y) + self.sig(y) + + +class DiamondChained(nn.Module): + """Same node count and node types as DiamondShared but different connectivity.""" + + def __init__(self): + super().__init__() + self.fc = nn.Linear(4, 4) + self.relu = nn.ReLU() + self.sig = nn.Sigmoid() + + def forward(self, x): + y = self.fc(x) + r = self.relu(y) + return self.sig(r) + r + + +class AddTwoInputs(nn.Module): + def forward(self, x, y): + return x + y + + +class ReluUnusedInput(nn.Module): + """Same node count as AddTwoInputs but the output node has a single input.""" + + def forward(self, x, y): + return F.relu(x) + + +class TwoOutputs(nn.Module): + def forward(self, x): + return x.relu(), x.sigmoid() + + +class ChainedOutput(nn.Module): + """Same node count and node types as TwoOutputs but a single chained output.""" + + def forward(self, x): + return x.relu().sigmoid() + + +class Untraceable(nn.Module): + """Data-dependent control flow makes this module untraceable by torch.fx.""" + + def forward(self, x): + if x.sum() > 0: + return x + return -x + + +@pytest.mark.parametrize( + ("make_module", "make_pattern"), + [ + (LinearRelu, LinearRelu), + # different layer sizes must still match since submodules are compared by type + (lambda: LinearRelu(features=4), lambda: LinearRelu(features=16)), + (LinearFuncRelu, LinearFuncRelu), + (LinearMethodRelu, LinearMethodRelu), + # equivalent graph built from different module classes + (lambda: nn.Sequential(nn.Linear(8, 8), nn.ReLU()), LinearRelu), + # shared intermediate node is matched consistently on revisit + (DiamondShared, DiamondShared), + ], +) +def test_match_equivalent_graphs(make_module, make_pattern): + assert match(make_module(), [make_pattern()]) + + +@pytest.mark.parametrize( + ("make_module", "make_pattern"), + [ + (LinearRelu, LinearSigmoid), # different call_module target type + (JustLinear, LinearRelu), # different node count + (LinearRelu, JustLinear), + (LinearRelu, LinearFuncRelu), # call_module vs call_function op + (LinearMethodRelu, LinearFuncRelu), # call_method vs call_function op + (AddTwoInputs, ReluUnusedInput), # same node count, different input arity + (DiamondShared, DiamondChained), # same node count/types, different connectivity + (TwoOutputs, ChainedOutput), # output nodes differ in input degree + ], +) +def test_no_match(make_module, make_pattern): + assert not match(make_module(), [make_pattern()]) + + +def test_empty_patterns_never_match(): + assert not match(LinearRelu(), []) + + +def test_match_any_of_multiple_patterns(): + assert match(LinearRelu(), [LinearSigmoid(), JustLinear(), LinearRelu()]) + + +def test_untraceable_module_returns_false(): + assert not match(Untraceable(), [LinearRelu()]) diff --git a/tests/unit/torch/utils/test_logging.py b/tests/unit/torch/utils/test_logging.py new file mode 100644 index 00000000000..2a55a6b0718 --- /dev/null +++ b/tests/unit/torch/utils/test_logging.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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 sys +import warnings + +import pytest +import tqdm + +from modelopt.torch.utils.logging import ( + capture_io, + no_stdout, + silence_matched_warnings, + warn_rank_0, +) + + +def test_no_stdout_disables_tqdm_and_restores_it(): + with no_stdout(): + bar = tqdm.tqdm(range(2)) + assert bar.disable + bar.close() + bar2 = tqdm.tqdm(range(2)) + assert not bar2.disable # original tqdm.__init__ restored on exit + bar2.close() + + +def test_warn_rank_0_wraps_message_in_yellow_on_tty(monkeypatch): + monkeypatch.setattr(sys.stderr, "isatty", lambda: True) + with pytest.warns(UserWarning) as record: + warn_rank_0("colored message") + assert str(record[0].message) == "\033[33mcolored message\033[0m" + + +def test_capture_io_captures_stdout_and_stderr(): + with capture_io() as buf: + print("to stdout") + print("to stderr", file=sys.stderr) + assert "to stdout" in buf.getvalue() + assert "to stderr" in buf.getvalue() + + +def test_capture_io_can_leave_stderr_alone(capsys): + with capture_io(capture_stderr=False) as buf: + print("out") + print("err", file=sys.stderr) + assert buf.getvalue() == "out\n" + assert "err" in capsys.readouterr().err + + +def test_silence_matched_warnings_filters_by_pattern(): + seen = [] + + def recorder(message, category, filename, lineno, file=None, line=None): + seen.append(str(message)) + + with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.showwarning = recorder # restored by catch_warnings on exit + with silence_matched_warnings("skip"): + warnings.warn("please skip me") + warnings.warn("keep me") + assert seen == ["keep me"] + + +def test_silence_matched_warnings_restores_showwarning(): + original = warnings.showwarning + with silence_matched_warnings("pattern"): + assert warnings.showwarning is not original + assert warnings.showwarning is original + + +@pytest.mark.parametrize("pattern", [None, 123]) +def test_silence_matched_warnings_invalid_pattern_is_noop(pattern): + # None or an un-compilable pattern leaves warnings.showwarning untouched + original = warnings.showwarning + with silence_matched_warnings(pattern): + assert warnings.showwarning is original + assert warnings.showwarning is original diff --git a/tests/unit/torch/utils/test_robust_json.py b/tests/unit/torch/utils/test_robust_json.py new file mode 100644 index 00000000000..47e83426c2b --- /dev/null +++ b/tests/unit/torch/utils/test_robust_json.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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 argparse +import datetime +import json +from enum import Enum + +import pytest +import torch + +from modelopt.torch.utils.robust_json import json_dumps + + +class _Color(Enum): + RED = "red" + GREEN = "green" + + +@pytest.mark.parametrize( + ("obj", "expected"), + [ + (_Color.RED, "RED"), # Enum encoded as name, not value + (argparse.Namespace(lr=0.1, name="exp"), {"lr": 0.1, "name": "exp"}), + (torch.bfloat16, "torch.bfloat16"), # torch dtype encoded as string + (datetime.timedelta(hours=1, minutes=2, seconds=3), "1:02:03"), + ], +) +def test_special_types_encoded(obj, expected): + assert json.loads(json_dumps({"k": obj})) == {"k": expected} + + +def test_main_module_function_encoded_as_bare_name(): + def user_fn(): + pass + + # user-defined functions in __main__ fall back to just the name + user_fn.__module__ = "__main__" + assert json.loads(json_dumps({"f": user_fn})) == {"f": "user_fn"}