rebase bug fixes

Signed-off-by: Kinjal Patel <[email protected]>

Author: kinjalpatel27

tests/_test_utils/torch_quantization/quantize_common.py

@@ -210,103 +210,6 @@ def forward_loop(model):
         )
 
 
-@patch("modelopt.torch.quantization.model_calib.awq_lite", side_effect=_debug_awq_lite)
-def dp_cp_parallel_test_helper(model, config, group, mock_awq_lite):
-    calib_data = model.get_dummy_input().cuda()
-
-    def forward_loop(model):
-        model(calib_data)
-
-    model = mtq.quantize(model, config, forward_loop)
-
-    # Sanity check
-    forward_loop(model)
-
-    # Input quantizer amax
-    if config not in [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT4_AWQ_CFG]:
-        _reduce_quantizer_attr(model.fc1.input_quantizer, "amax", dist.ReduceOp.MAX, group=group)
-        _reduce_quantizer_attr(model.fc2.input_quantizer, "amax", dist.ReduceOp.MAX, group=group)
-
-    # Weight quantizer amax
-    if isinstance(model.fc1.weight_quantizer, SequentialQuantizer):
-        for quantizer in model.fc1.weight_quantizer:
-            _reduce_quantizer_attr(quantizer, "amax", dist.ReduceOp.MAX, group=group)
-    else:
-        _reduce_quantizer_attr(model.fc1.weight_quantizer, "amax", dist.ReduceOp.MAX, group=group)
-    if isinstance(model.fc2.weight_quantizer, SequentialQuantizer):
-        for quantizer in model.fc2.weight_quantizer:
-            _reduce_quantizer_attr(quantizer, "amax", dist.ReduceOp.MAX, group=group)
-    else:
-        _reduce_quantizer_attr(model.fc2.weight_quantizer, "amax", dist.ReduceOp.MAX, group=group)
-
-    if config in [mtq.INT4_AWQ_CFG, mtq.W4A8_AWQ_BETA_CFG]:
-        # Check act scale
-        _reduce_quantizer_attr(
-            model.fc1.awq_lite,
-            "act_scale",
-            dist.ReduceOp.AVG,
-            group=group,
-        )
-        _reduce_quantizer_attr(
-            model.fc2.awq_lite,
-            "act_scale",
-            dist.ReduceOp.AVG,
-            group=group,
-        )
-
-
-@patch("modelopt.torch.quantization.model_calib.awq_lite", side_effect=_debug_awq_lite)
-def data_tensor_context_parallel_test_helper(model, config, dp_group, tp_group, mock_awq_lite):
-    # Calib data should be same across each DP rank
-    dp_rank = dist.get_rank(group=dp_group)
-    calib_data = model.get_dummy_input(seed=dp_rank).cuda()
-
-    def forward_loop(model):
-        model(calib_data)
-
-    model = mtq.quantize(model, config, forward_loop)
-
-    def _reduce_quantizer_attr(quantizer, attr=str, op=dist.ReduceOp.MAX):
-        quantizer_attr = getattr(quantizer, attr).clone()
-
-        # Perform all-reduce operations
-        dist.all_reduce(quantizer_attr, op=op, group=tp_group)
-
-        dist.all_reduce(quantizer_attr, op=op, group=dp_group)
-
-        assert torch.allclose(quantizer_attr, getattr(quantizer, attr)), getattr(quantizer, attr)
-
-    # Input quantizer amax
-    if config not in [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT4_AWQ_CFG]:
-        _reduce_quantizer_attr(model.fc1.input_quantizer, "amax", dist.ReduceOp.MAX)
-        _reduce_quantizer_attr(model.fc2.input_quantizer, "amax", dist.ReduceOp.MAX)
-
-    # Per-tensor quantization (FP8/NVFP4) expects same amax across row and column parallel ranks
-    # Channel-wise (INT8) only expects same amax across row parallel ranks
-    # Block-wise quantization does not expect same amax across row and column parallel ranks
-    if config in [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]:
-        if isinstance(model.fc1.weight_quantizer, SequentialQuantizer):
-            for quantizer in model.fc1.weight_quantizer:
-                _reduce_quantizer_attr(quantizer, "amax", dist.ReduceOp.MAX)
-        else:
-            _reduce_quantizer_attr(model.fc1.weight_quantizer, "amax", dist.ReduceOp.MAX)
-
-    if config in [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG, mtq.INT8_DEFAULT_CFG]:
-        if isinstance(model.fc2.weight_quantizer, SequentialQuantizer):
-            for quantizer in model.fc2.weight_quantizer:
-                _reduce_quantizer_attr(quantizer, "amax", dist.ReduceOp.MAX)
-        else:
-            _reduce_quantizer_attr(model.fc2.weight_quantizer, "amax", dist.ReduceOp.MAX)
-
-    # Check act scale
-    if config in [mtq.INT4_AWQ_CFG, mtq.W4A8_AWQ_BETA_CFG]:
-        _reduce_quantizer_attr(
-            model.fc1.awq_lite,
-            "act_scale",
-            dist.ReduceOp.AVG,
-        )
-
-
 def auto_quantize_helper(model):
     model, search_state = mtq.auto_quantize(
         model,

tests/gpu/torch/quantization/plugins/test_megatron.py

@@ -231,89 +231,6 @@ def test_data_tensor_context_parallel(need_8_gpus, config):
     )
 
 
-# 2. Data Parallel Test
-def _test_data_parallel_helper(config, rank, size):
-    initialize_for_megatron(seed=SEED + rank)  # modify seed so data is different across ranks
-    model = MegatronModel().cuda()
-
-    dp_cp_parallel_test_helper(model, config, get_data_parallel_group())
-
-
-@pytest.mark.parametrize(
-    "config",
-    [
-        mtq.INT8_DEFAULT_CFG,
-        mtq.FP8_DEFAULT_CFG,
-        mtq.W4A8_AWQ_BETA_CFG,
-        mtq.INT8_SMOOTHQUANT_CFG,
-        mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
-        mtq.INT4_AWQ_CFG,
-        mtq.NVFP4_DEFAULT_CFG,
-    ],
-)
-def test_data_parallel(need_2_gpus, config):
-    spawn_multiprocess_job(size=2, job=partial(_test_data_parallel_helper, config), backend="nccl")
-
-
-# 3. Context Parallel Test
-def _test_context_parallel_helper(config, rank, size):
-    initialize_for_megatron(
-        context_parallel_size=size, seed=SEED + rank
-    )  # modify seed so data is different across ranks
-    model = MegatronModel(cp_size=size).cuda()
-
-    dp_cp_parallel_test_helper(model, config, get_data_parallel_group(with_context_parallel=True))
-
-
-@pytest.mark.parametrize(
-    "config",
-    [
-        mtq.INT8_DEFAULT_CFG,
-        mtq.FP8_DEFAULT_CFG,
-        mtq.W4A8_AWQ_BETA_CFG,
-        mtq.INT8_SMOOTHQUANT_CFG,
-        mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
-        mtq.INT4_AWQ_CFG,
-        mtq.NVFP4_DEFAULT_CFG,
-    ],
-)
-def test_context_parallel(need_2_gpus, config):
-    spawn_multiprocess_job(
-        size=2, job=partial(_test_context_parallel_helper, config), backend="nccl"
-    )
-
-
-# 4. DP=2 + TP=2 + CP=2 Test (on 2*2*2=8 GPUs)
-def _test_data_tensor_context_parallel_helper(config, rank, size):
-    initialize_for_megatron(tensor_model_parallel_size=2, context_parallel_size=2, seed=SEED + rank)
-    model = MegatronModel(tp_size=2, cp_size=2).cuda()
-
-    data_tensor_context_parallel_test_helper(
-        model,
-        config,
-        get_data_parallel_group(with_context_parallel=True),
-        get_tensor_model_parallel_group(),
-    )
-
-
-@pytest.mark.parametrize(
-    "config",
-    [
-        mtq.INT8_DEFAULT_CFG,
-        mtq.FP8_DEFAULT_CFG,
-        mtq.W4A8_AWQ_BETA_CFG,
-        mtq.INT8_SMOOTHQUANT_CFG,
-        mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
-        mtq.INT4_AWQ_CFG,
-        mtq.NVFP4_DEFAULT_CFG,
-    ],
-)
-def test_data_tensor_context_parallel(need_8_gpus, config):
-    spawn_multiprocess_job(
-        size=8, job=partial(_test_data_tensor_context_parallel_helper, config), backend="nccl"
-    )
-
-
 def _gpt_model_provider(
     tp_size: int,
     hidden_size=256,