diff --git a/test/optimizations/test_bevfusion_implementations.py b/test/optimizations/test_bevfusion_implementations.py index 4866b6c..daeb3ff 100644 --- a/test/optimizations/test_bevfusion_implementations.py +++ b/test/optimizations/test_bevfusion_implementations.py @@ -18,6 +18,62 @@ from turbo_physai.optimizations.models.bevfusion import transfusion +def test_ddp_forward_supports_new_pytorch(): + class DDPStub(torch.nn.Module): + def __init__(self): + super().__init__() + self.module = torch.nn.Linear(2, 1) + + def _run_ddp_forward(self, *inputs, **kwargs): + assert self._use_replicated_tensor_module is False + return self.module(*inputs, **kwargs) + + model = DDPStub() + wrapped = training.ddp_forward_compat_wrapper( + DDPStub._run_ddp_forward, {} + ) + + assert wrapped(model, torch.ones(1, 2)).shape == (1, 1) + assert model._use_replicated_tensor_module is False + + +@pytest.mark.parametrize("compiled", [False, True]) +def test_bev_pool_fp16_casts_factorized_namedtuples(monkeypatch, compiled): + def pool(self, geometry, factors): + return geometry, factors + + monkeypatch.setattr(depth, "base_transform_bev_pool", pool) + wrapped = depth.base_transform_bev_pool_wrapper(None, {}) + if compiled: + wrapped = torch.compile(wrapped, backend="eager") + + model = torch.nn.Module() + model.fp16_enabled = True + geometry = depth.PreparedGeometry( + torch.ones(1, dtype=torch.float16), + torch.ones(1, dtype=torch.long), + torch.ones(1, dtype=torch.bool), + 1, + ) + factors = depth.DepthFeatureFactorization( + torch.ones(1, dtype=torch.float16), + torch.ones(1, dtype=torch.float16), + ) + actual_geometry, actual_factors = wrapped(model, geometry, factors) + assert isinstance(actual_geometry, depth.PreparedGeometry) + assert isinstance(actual_factors, depth.DepthFeatureFactorization) + assert actual_geometry.coords.dtype == torch.float32 + assert actual_geometry.ranks.dtype == torch.long + assert actual_geometry.kept.dtype == torch.bool + assert actual_factors.depth.dtype == torch.float32 + assert actual_factors.features.dtype == torch.float32 + + model.fp16_enabled = False + unchanged_geometry, unchanged_factors = wrapped(model, geometry, factors) + assert unchanged_geometry is geometry + assert unchanged_factors is factors + + def test_extract_camera_features_preserves_camera_contract(monkeypatch): captured = {} diff --git a/turbo_physai/optimizations/models/bevfusion/catalog.py b/turbo_physai/optimizations/models/bevfusion/catalog.py index fe10815..1d8c07f 100644 --- a/turbo_physai/optimizations/models/bevfusion/catalog.py +++ b/turbo_physai/optimizations/models/bevfusion/catalog.py @@ -88,6 +88,13 @@ "turbo_physai.optimizations.models.bevfusion.training.training_wrapper" ), ), + wrap( + target="mmcv.parallel.distributed.MMDistributedDataParallel._run_ddp_forward", + replacement=( + "turbo_physai.optimizations.models.bevfusion.training." + "ddp_forward_compat_wrapper" + ), + ), ) GAUSSIAN = group( diff --git a/turbo_physai/optimizations/models/bevfusion/configs/.optimization.yaml.generation.json b/turbo_physai/optimizations/models/bevfusion/configs/.optimization.yaml.generation.json index 3718069..5fe95a7 100644 --- a/turbo_physai/optimizations/models/bevfusion/configs/.optimization.yaml.generation.json +++ b/turbo_physai/optimizations/models/bevfusion/configs/.optimization.yaml.generation.json @@ -4,7 +4,7 @@ "model_commit": "326653dc06e0938edf1aae7d01efcd158ba83de5", "config": { "path": "optimization.yaml", - "sha256": "eadc1cf9155cf8279fad9fdae9188d24b9e136edecf6abb920175fd71ba4aed1" + "sha256": "ec41d78a5c8c5b8c053f47bb1996ddab72ee683e647902f71d07b951d13baf2d" }, "inputs": [ { @@ -35,7 +35,7 @@ "role": "catalog", "base": "package", "path": "optimizations/models/bevfusion/catalog.py", - "sha256": "309abcd7ea5c9f7cd64a89fab30ded23b62922397e711018fa155f5af9f88e21" + "sha256": "8fdfa085391af2e3bec60deaab92d1f66f0195d517c5f4877b55156eb184441c" } ] } diff --git a/turbo_physai/optimizations/models/bevfusion/configs/optimization.yaml b/turbo_physai/optimizations/models/bevfusion/configs/optimization.yaml index 5d92dd3..c29a09b 100644 --- a/turbo_physai/optimizations/models/bevfusion/configs/optimization.yaml +++ b/turbo_physai/optimizations/models/bevfusion/configs/optimization.yaml @@ -124,9 +124,13 @@ optimization_groups: source_hashes: mmdet3d.apis.train.train_model: - source-v1:6215bc796d762e6fea35f7289bf6edc46f9f406e377520e00bc7a500144080b1 + mmcv.parallel.distributed.MMDistributedDataParallel._run_ddp_forward: + - source-v1:f4e196a85655850dbae1bf752ea1519a86083ba518b95a34cc5afbcf76bb5274 ast_hashes: mmdet3d.apis.train.train_model: - ast-v1:b1e7df6eca71976ac570963e0ee3419fa294faa6501a9c9321613779442bf2a7 + mmcv.parallel.distributed.MMDistributedDataParallel._run_ddp_forward: + - ast-v1:7d601b5fefa14e6288d8b597d1b68043025af1567dc704226bdd24c17faab5d9 - id: bevfusion.gaussian enabled: true options: {} diff --git a/turbo_physai/optimizations/models/bevfusion/depth.py b/turbo_physai/optimizations/models/bevfusion/depth.py index 6db6301..71a77e4 100644 --- a/turbo_physai/optimizations/models/bevfusion/depth.py +++ b/turbo_physai/optimizations/models/bevfusion/depth.py @@ -641,12 +641,40 @@ def base_transform_get_geometry( def base_transform_bev_pool_wrapper(original, options): - """Retain BaseTransform's MMCV FP32 boundary around optimized pooling.""" + """Cast pooling inputs to FP32 without MMCV's broken NamedTuple recursion.""" del original, options - from mmcv.runner import force_fp32 + import functools + import torch - return force_fp32()(base_transform_bev_pool) + def fp32(value): + if isinstance(value, torch.Tensor) and value.dtype == torch.float16: + return value.float() + return value + + @functools.wraps(base_transform_bev_pool) + def wrapped(self, geom_feats, x): + if not getattr(self, "fp16_enabled", False): + return base_transform_bev_pool(self, geom_feats, x) + + if isinstance(geom_feats, PreparedGeometry): + geom_feats = PreparedGeometry( + fp32(geom_feats.coords), + geom_feats.ranks, + geom_feats.kept, + geom_feats.batch_size, + ) + else: + geom_feats = fp32(geom_feats) + if isinstance(x, DepthFeatureFactorization): + x = DepthFeatureFactorization(fp32(x.depth), fp32(x.features)) + else: + x = fp32(x) + + with torch.amp.autocast("cuda", enabled=False): + return base_transform_bev_pool(self, geom_feats, x) + + return wrapped def base_transform_bev_pool_prepared(self, x, coords, ranks, kept): diff --git a/turbo_physai/optimizations/models/bevfusion/training.py b/turbo_physai/optimizations/models/bevfusion/training.py index 5af90c3..c22e7af 100644 --- a/turbo_physai/optimizations/models/bevfusion/training.py +++ b/turbo_physai/optimizations/models/bevfusion/training.py @@ -36,6 +36,20 @@ def _option_flag(options, name, env_name, default=False): return _env_flag(env_name, default) +def ddp_forward_compat_wrapper(original, options): + """Support MMCV's DDP forward on newer PyTorch versions.""" + + del options + + @functools.wraps(original) + def wrapped(self, *args, **kwargs): + if not hasattr(self, "_use_replicated_tensor_module"): + self._use_replicated_tensor_module = False + return original(self, *args, **kwargs) + + return wrapped + + def parse_losses(self, losses): """Reduce all scalar losses with one collective and one host transfer."""