lint

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

Author: jenchen13

modelopt/torch/quantization/model_calib.py

@@ -628,10 +628,14 @@ def forward(self, input, *args, **kwargs):
     def sync_act_scale_across_dp_cp(module, data_parallel_group, context_parallel_group):
         # Sync across Data Parallel (DP)
         if data_parallel_group.is_initialized():
-            dist.all_reduce(module.awq_lite.act_scale, op=dist.ReduceOp.AVG, group=data_parallel_group.group)
+            dist.all_reduce(
+                module.awq_lite.act_scale, op=dist.ReduceOp.AVG, group=data_parallel_group.group
+            )
         # Sync across Context Parallel (CP)
         if context_parallel_group.is_initialized():
-            dist.all_reduce(module.awq_lite.act_scale, op=dist.ReduceOp.AVG, group=context_parallel_group.group)
+            dist.all_reduce(
+                module.awq_lite.act_scale, op=dist.ReduceOp.AVG, group=context_parallel_group.group
+            )
 
     for name, module in model.named_modules():
         if (
@@ -640,8 +644,12 @@ def sync_act_scale_across_dp_cp(module, data_parallel_group, context_parallel_gr
             and module.awq_lite.num_cache_steps > 0
         ):
             module.awq_lite.act_scale = module.awq_lite.act_scale / module.awq_lite.num_cache_steps
-            sync_act_scale_across_dp_cp(module, module.parallel_state.data_parallel_group, module.parallel_state.context_parallel_group)
-            
+            sync_act_scale_across_dp_cp(
+                module,
+                module.parallel_state.data_parallel_group,
+                module.parallel_state.context_parallel_group,
+            )
+
             # Hack: MoEs forward all tokens through all experts if _if_calib is True
             module._if_calib = True
 

modelopt/torch/quantization/plugins/megatron.py

@@ -22,8 +22,8 @@
 import megatron.core.tensor_parallel.layers as megatron_parallel
 import megatron.core.transformer.mlp as megatron_mlp
 import torch
-from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
 from megatron.core.parallel_state import get_data_parallel_group
+from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
 from megatron.core.transformer import MegatronModule
 from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
 from megatron.core.utils import get_tensor_model_parallel_group_if_none
@@ -221,7 +221,7 @@ def _setup(self):
         data_parallel_group = None
         try:
             data_parallel_group = get_data_parallel_group(with_context_parallel=True)
-        except:
+        except AssertionError:
             data_parallel_group = get_data_parallel_group()
         self.parallel_state = ParallelState(
             data_parallel_group,

modelopt/torch/utils/distributed.py

@@ -249,7 +249,11 @@ def __init__(
         self.context_parallel_group = DistributedProcessGroup(context_parallel_group)
 
     def __repr__(self) -> str:
-        return f"data_parallel_group: {self.data_parallel_group}, tensor_parallel_group: {self.tensor_parallel_group}, context_parallel_group: {self.context_parallel_group}"
+        return (
+            f"data_parallel_group: {self.data_parallel_group}, "
+            f"tensor_parallel_group: {self.tensor_parallel_group}, "
+            f"context_parallel_group: {self.context_parallel_group}"
+        )
 
 
 def get_group(ranks: list[int]):

tests/_test_utils/torch_dist/plugins/megatron_common.py

@@ -390,7 +390,11 @@ def initialize_for_megatron(
 
     NOTE: If used in a non-spawned process, make sure to call `megatron.core.parallel_state.destroy_model_parallel()`.
     """
-    initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, context_parallel_size=context_parallel_size)
+    initialize_model_parallel(
+        tensor_model_parallel_size,
+        pipeline_model_parallel_size,
+        context_parallel_size=context_parallel_size,
+    )
     model_parallel_cuda_manual_seed(seed)
 
 

tests/_test_utils/torch_quantization/quantize_common.py

@@ -149,6 +149,7 @@ def forward_loop(model):
 
     dist.destroy_process_group()
 
+
 def data_parallel_test_helper(model, config, dp_group):
     calib_data = model.get_dummy_input().cuda()
 
@@ -165,6 +166,7 @@ def forward_loop(model):
     dist.all_reduce(fc2_amax, op=dist.ReduceOp.MAX, group=dp_group)
     assert torch.allclose(fc2_amax, model.fc2.input_quantizer.amax)
 
+
 def context_parallel_test_helper(model, config, cp_group):
     calib_data = model.get_dummy_input().cuda()
 
@@ -181,6 +183,7 @@ def forward_loop(model):
     dist.all_reduce(fc2_amax, op=dist.ReduceOp.MAX, group=cp_group)
     assert torch.allclose(fc2_amax, model.fc2.input_quantizer.amax)
 
+
 def data_tensor_context_parallel_test_helper(model, config, dp_group, tp_group, cp_group):
     calib_data = model.get_dummy_input().cuda()
     # data should be same across each TP rank
@@ -203,6 +206,7 @@ def forward_loop(model):
     dist.all_reduce(fc2_amax, op=dist.ReduceOp.MAX, group=dp_group)
     assert torch.allclose(fc2_amax, model.fc2.input_quantizer.amax)
 
+
 def auto_quantize_helper(model):
     model, search_state = mtq.auto_quantize(
         model,

tests/gpu/torch/conftest.py

@@ -33,13 +33,13 @@ def need_2_gpus():
     if torch.cuda.device_count() < 2:
         pytest.skip("Need at least 2 GPUs to run this test")
 
+
 @pytest.fixture
 def need_8_gpus():
     if torch.cuda.device_count() < 8:
         pytest.skip("Need at least 8 GPUs to run this test")
 
 
-
 @pytest.fixture(scope="module")
 def set_torch_dtype(request):
     orig_dtype = torch.get_default_dtype()

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

@@ -31,10 +31,10 @@
 from _test_utils.torch_quantization.quant_utils import get_model_size
 from _test_utils.torch_quantization.quantize_common import (
     auto_quantize_helper,
-    tensor_parallel_test_helper,
-    data_parallel_test_helper,
     context_parallel_test_helper,
+    data_parallel_test_helper,
     data_tensor_context_parallel_test_helper,
+    tensor_parallel_test_helper,
 )
 from packaging.version import Version
 
@@ -43,8 +43,8 @@
 import megatron.core
 from megatron.core.parallel_state import (
     destroy_model_parallel,
-    get_data_parallel_group,
     get_context_parallel_group,
+    get_data_parallel_group,
     get_tensor_model_parallel_group,
 )
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
@@ -95,14 +95,13 @@ def test_convert_megatron_parallel_linear(distributed_setup_size_1):
     # Clean up since this is not a spawned process
     destroy_model_parallel()
 
+
 # 1. Tensor Parallel Test
 def _test_tensor_parallel_helper(config, rank, size):
     initialize_for_megatron(tensor_model_parallel_size=2, seed=SEED)
     model = MegatronModel(tp_size=size).cuda()
 
-    tensor_parallel_test_helper(
-        model, config, get_tensor_model_parallel_group()
-    )
+    tensor_parallel_test_helper(model, config, get_tensor_model_parallel_group())
 
 
 @pytest.mark.parametrize(
@@ -122,15 +121,14 @@ def test_tensor_parallel(need_2_gpus, config):
         size=2, job=partial(_test_tensor_parallel_helper, config), backend="nccl"
     )
 
+
 # 2. Data Parallel Test
 def _test_data_parallel_helper(config, rank, size):
     # TODO does this model automatically get copied to both DP ranks?
     initialize_for_megatron(seed=SEED)
     model = MegatronModel().cuda()
 
-    data_parallel_test_helper(
-        model, config, get_data_parallel_group()
-    )
+    data_parallel_test_helper(model, config, get_data_parallel_group())
 
 
 @pytest.mark.parametrize(
@@ -146,18 +144,16 @@ def _test_data_parallel_helper(config, rank, size):
     ],
 )
 def test_data_parallel(need_2_gpus, config):
-    spawn_multiprocess_job(
-        size=2, job=partial(_test_data_parallel_helper, config), backend="nccl"
-    )
+    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)
     model = MegatronModel(cp_size=size).cuda()
 
-    context_parallel_test_helper(
-        model, config, get_context_parallel_group()
-    )
+    context_parallel_test_helper(model, config, get_context_parallel_group())
+
 
 @pytest.mark.parametrize(
     "config",
@@ -176,15 +172,21 @@ def test_context_parallel(need_2_gpus, config):
         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)
     model = MegatronModel(tp_size=2, cp_size=2).cuda()
 
     data_tensor_context_parallel_test_helper(
-        model, config, get_data_parallel_group(), get_tensor_model_parallel_group(), get_context_parallel_group()
+        model,
+        config,
+        get_data_parallel_group(),
+        get_tensor_model_parallel_group(),
+        get_context_parallel_group(),
     )
 
+
 @pytest.mark.parametrize(
     "config",
     [
@@ -199,9 +201,10 @@ def _test_data_tensor_context_parallel_helper(config, rank, size):
 )
 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"
+        size=8, job=partial(_test_data_tensor_context_parallel_helper, config), backend="nccl"
     )
 
+
 def _gpt_model_provider(tp_size: int, hidden_size=256, vocab_size=64, meta_device=False):
     """Build the model."""