lint

Signed-off-by: wdykas <wdykas@nvidia.com>

Author: wdykas

megatron/core/inference/shards.py

@@ -75,9 +75,7 @@ def build_inference_pg_collection(
 
     # Expert DP size (world = expt_tp * ep * expt_dp * pp)
     expt_dp_size = world_size // (expt_tp_size * ep_size * pp_size)
-    assert expt_dp_size >= 1 and (
-        expt_tp_size * ep_size * expt_dp_size * pp_size
-    ) == world_size, (
+    assert expt_dp_size >= 1 and (expt_tp_size * ep_size * expt_dp_size * pp_size) == world_size, (
         f"World size ({world_size}) must be divisible by expt_tp*ep*pp "
         f"({expt_tp_size * ep_size * pp_size})"
     )
@@ -86,15 +84,11 @@ def build_inference_pg_collection(
 
     if use_tp_pp_dp_mapping:
         decoder_grid = HyperCommGrid(
-            [tp_size, cp_size, pp_size, dp_size],
-            ["tp", "cp", "pp", "dp"],
-            rank_offset=rank_offset,
+            [tp_size, cp_size, pp_size, dp_size], ["tp", "cp", "pp", "dp"], rank_offset=rank_offset
         )
     else:
         decoder_grid = HyperCommGrid(
-            [tp_size, cp_size, dp_size, pp_size],
-            ["tp", "cp", "dp", "pp"],
-            rank_offset=rank_offset,
+            [tp_size, cp_size, dp_size, pp_size], ["tp", "cp", "dp", "pp"], rank_offset=rank_offset
         )
 
     tp_group = decoder_grid.create_pg("tp")

tests/unit_tests/inference/test_inference_shards.py

@@ -12,8 +12,7 @@
 
 
 def test_shard_spec_objects_match_string_parsing():
-    objs = [InferenceShardSpec(tp=2, role="prefill"),
-            InferenceShardSpec(tp=1, dp=2, role="decode")]
+    objs = [InferenceShardSpec(tp=2, role="prefill"), InferenceShardSpec(tp=1, dp=2, role="decode")]
     assert normalize_shard_specs(objs, 4) == parse_inference_shards_spec(
         "tp=2,role=prefill+tp=1,dp=2,role=decode", 4
     )
@@ -34,7 +33,14 @@ def test_parse_defaults_and_dp_and_role():
     assert specs[0] == InferenceShardSpec(tp=2, role="prefill")
     assert specs[1] == InferenceShardSpec(tp=1, dp=2, role="decode")
     # dict form (serialization / external consumers) carries the resolved keys.
-    assert specs[0].to_dict() == {"tp": 2, "pp": 1, "ep": 1, "dp": 1, "expt_tp": 2, "role": "prefill"}
+    assert specs[0].to_dict() == {
+        "tp": 2,
+        "pp": 1,
+        "ep": 1,
+        "dp": 1,
+        "expt_tp": 2,
+        "role": "prefill",
+    }
 
 
 def test_parse_partitions_world_with_dp():
@@ -60,7 +66,9 @@ def test_plus_and_semicolon_separators_equivalent():
 def test_cp_accepted_only_when_one():
     # cp is a recognized key (clear error, not "unknown key") but must be 1:
     # inference shards don't context-parallelize.
-    assert parse_inference_shards_spec("tp=2,cp=1", world_size=2) == [InferenceShardSpec(tp=2, cp=1)]
+    assert parse_inference_shards_spec("tp=2,cp=1", world_size=2) == [
+        InferenceShardSpec(tp=2, cp=1)
+    ]
     with pytest.raises(ValueError):
         InferenceShardSpec(tp=1, cp=2)
     with pytest.raises(ValueError):