Fixed

Signed-off-by: greg-kwasniewski1 <[email protected]>

Author: greg-kwasniewski1

setup.py

@@ -215,7 +215,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
             precompiled_location = download_precompiled(tempdir, version)
         extract_from_precompiled(precompiled_location, package_data, tempdir)
 
-# sanity_check()
+sanity_check()
 
 # https://setuptools.pypa.io/en/latest/references/keywords.html
 setup(

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -56,7 +56,7 @@ transforms:
     stage: sharding
     simple_shard_only: false
     use_sharding_from_factory: false
-    sharding_dims: ['tp', 'ep', 'dp']
+    sharding_dims: ['tp', 'ep', 'bmm']
   # TODO: (hg) need to ensure run_shape_prop after sharding.
   sharding_transform_executor:
     stage: sharding

tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py

@@ -67,6 +67,7 @@ def _run_job(
             "detect_sharding": {
                 "stage": "sharding",
                 "use_sharding_from_factory": False,
+                "sharding_dims": ["bmm"],
             },
             "sharding_transform_executor": {
                 "stage": "sharding",