|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | +# All rights reserved. |
| 4 | +# |
| 5 | +# This source code is licensed under the BSD-style license found in the |
| 6 | +# LICENSE file in the root directory of this source tree. |
| 7 | + |
| 8 | +# pyre-strict |
| 9 | + |
| 10 | +# This test exercises DATA_PARALLEL replication and the DDP-wrap decision, which |
| 11 | +# are device-independent; run it CPU/gloo-only so it needs no GPU. CUDA is hidden |
| 12 | +# per-test (see the class below) to keep MultiProcessContext on its CPU path |
| 13 | +# without mutating the process-wide env at import time. |
| 14 | +import os |
| 15 | +import unittest |
| 16 | +from typing import cast, List |
| 17 | +from unittest.mock import patch |
| 18 | + |
| 19 | +import torch |
| 20 | +from torch import nn |
| 21 | +from torch.nn.parallel import DistributedDataParallel |
| 22 | +from torchrec.distributed.embedding import ( |
| 23 | + EmbeddingCollectionSharder, |
| 24 | + ShardedEmbeddingCollection, |
| 25 | +) |
| 26 | +from torchrec.distributed.embeddingbag import ( |
| 27 | + EmbeddingBagCollectionSharder, |
| 28 | + ShardedEmbeddingBagCollection, |
| 29 | +) |
| 30 | +from torchrec.distributed.model_parallel import DistributedModelParallel |
| 31 | +from torchrec.distributed.sharding_plan import ( |
| 32 | + construct_module_sharding_plan, |
| 33 | + data_parallel, |
| 34 | +) |
| 35 | +from torchrec.distributed.test_utils.multi_process import ( |
| 36 | + MultiProcessContext, |
| 37 | + MultiProcessTestBase, |
| 38 | +) |
| 39 | +from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan |
| 40 | +from torchrec.modules.embedding_configs import EmbeddingBagConfig, EmbeddingConfig |
| 41 | +from torchrec.modules.embedding_modules import ( |
| 42 | + EmbeddingBagCollection, |
| 43 | + EmbeddingCollection, |
| 44 | + mark_data_parallel_skip_grad_sync, |
| 45 | + should_skip_data_parallel_grad_sync, |
| 46 | +) |
| 47 | +from torchrec.sparse.jagged_tensor import KeyedJaggedTensor |
| 48 | +from torchrec.test_utils import skip_if_asan_class |
| 49 | + |
| 50 | +_TABLE_NAME = "table_0" |
| 51 | +_FEATURE_NAME = "feature_0" |
| 52 | +_EMBEDDING_DIM = 8 |
| 53 | +_NUM_EMBEDDINGS = 16 |
| 54 | + |
| 55 | + |
| 56 | +def _build_ebc(device: torch.device) -> EmbeddingBagCollection: |
| 57 | + return EmbeddingBagCollection( |
| 58 | + tables=[ |
| 59 | + EmbeddingBagConfig( |
| 60 | + name=_TABLE_NAME, |
| 61 | + embedding_dim=_EMBEDDING_DIM, |
| 62 | + num_embeddings=_NUM_EMBEDDINGS, |
| 63 | + feature_names=[_FEATURE_NAME], |
| 64 | + ) |
| 65 | + ], |
| 66 | + device=device, |
| 67 | + ) |
| 68 | + |
| 69 | + |
| 70 | +def _build_ec(device: torch.device) -> EmbeddingCollection: |
| 71 | + return EmbeddingCollection( |
| 72 | + tables=[ |
| 73 | + EmbeddingConfig( |
| 74 | + name=_TABLE_NAME, |
| 75 | + embedding_dim=_EMBEDDING_DIM, |
| 76 | + num_embeddings=_NUM_EMBEDDINGS, |
| 77 | + feature_names=[_FEATURE_NAME], |
| 78 | + ) |
| 79 | + ], |
| 80 | + device=device, |
| 81 | + ) |
| 82 | + |
| 83 | + |
| 84 | +def _input_kjt(device: torch.device) -> KeyedJaggedTensor: |
| 85 | + return KeyedJaggedTensor( |
| 86 | + keys=[_FEATURE_NAME], |
| 87 | + values=torch.tensor([1, 2, 3], dtype=torch.long, device=device), |
| 88 | + lengths=torch.tensor([2, 1], dtype=torch.long, device=device), |
| 89 | + ) |
| 90 | + |
| 91 | + |
| 92 | +def _dp_lookups(dmp: DistributedModelParallel) -> List[nn.Module]: |
| 93 | + """Lookups of the FIRST sharded embedding module found in the DMP. Each DMP in |
| 94 | + this test wraps exactly one EBC/EC, so returning the first match is sufficient.""" |
| 95 | + for submodule in dmp.modules(): |
| 96 | + if isinstance( |
| 97 | + submodule, (ShardedEmbeddingBagCollection, ShardedEmbeddingCollection) |
| 98 | + ): |
| 99 | + # pyre-ignore[7]: private but stable attribute used for the assertion |
| 100 | + return list(submodule._lookups) |
| 101 | + raise AssertionError("no sharded embedding module found in DMP") |
| 102 | + |
| 103 | + |
| 104 | +def _resolve(output: object) -> object: |
| 105 | + return output.wait() if hasattr(output, "wait") else output |
| 106 | + |
| 107 | + |
| 108 | +def _shard_dp( |
| 109 | + module: nn.Module, |
| 110 | + sharder: ModuleSharder[nn.Module], |
| 111 | + ctx: MultiProcessContext, |
| 112 | + device: torch.device, |
| 113 | +) -> DistributedModelParallel: |
| 114 | + module_sharding_plan = construct_module_sharding_plan( |
| 115 | + module, |
| 116 | + per_param_sharding={_TABLE_NAME: data_parallel()}, |
| 117 | + local_size=ctx.local_size, |
| 118 | + world_size=ctx.world_size, |
| 119 | + device_type=device.type, |
| 120 | + ) |
| 121 | + return DistributedModelParallel( |
| 122 | + module=module, |
| 123 | + plan=ShardingPlan({"": module_sharding_plan}), |
| 124 | + # pyre-ignore[6]: ctx.pg is typed Optional[ProcessGroup]; MultiProcessContext |
| 125 | + # always initializes it before this call, so it is non-None here. |
| 126 | + env=ShardingEnv.from_process_group(ctx.pg), |
| 127 | + sharders=[sharder], |
| 128 | + device=device, |
| 129 | + ) |
| 130 | + |
| 131 | + |
| 132 | +def _run_dp_grad_sync_optout( |
| 133 | + rank: int, |
| 134 | + world_size: int, |
| 135 | + backend: str, |
| 136 | + is_ebc: bool, |
| 137 | +) -> None: |
| 138 | + with MultiProcessContext(rank, world_size, backend) as ctx: |
| 139 | + # Force CPU/gloo: DATA_PARALLEL replication and the DDP-wrap decision are |
| 140 | + # device-independent, and this keeps the test runnable without GPUs. |
| 141 | + device = torch.device("cpu") |
| 142 | + |
| 143 | + default_module = _build_ebc(device) if is_ebc else _build_ec(device) |
| 144 | + |
| 145 | + frozen_module = _build_ebc(device) if is_ebc else _build_ec(device) |
| 146 | + # Opt the frozen copy out of the DATA_PARALLEL DDP gradient reducer. |
| 147 | + mark_data_parallel_skip_grad_sync(frozen_module) |
| 148 | + |
| 149 | + sharder = cast( |
| 150 | + ModuleSharder[nn.Module], |
| 151 | + EmbeddingBagCollectionSharder() if is_ebc else EmbeddingCollectionSharder(), |
| 152 | + ) |
| 153 | + |
| 154 | + default_dmp = _shard_dp(default_module, sharder, ctx, device) |
| 155 | + frozen_dmp = _shard_dp(frozen_module, sharder, ctx, device) |
| 156 | + |
| 157 | + # Regression guard: the default DP module still gets a DDP reducer. |
| 158 | + default_lookups = _dp_lookups(default_dmp) |
| 159 | + assert any( |
| 160 | + isinstance(lookup, DistributedDataParallel) for lookup in default_lookups |
| 161 | + ), "default DATA_PARALLEL module must keep the DistributedDataParallel wrap" |
| 162 | + |
| 163 | + # The frozen/opt-out module has raw lookups (no DDP, no reducer/buckets). |
| 164 | + frozen_lookups = _dp_lookups(frozen_dmp) |
| 165 | + assert not any( |
| 166 | + isinstance(lookup, DistributedDataParallel) for lookup in frozen_lookups |
| 167 | + ), "opt-out DATA_PARALLEL module must skip the DistributedDataParallel wrap" |
| 168 | + |
| 169 | + # The unwrapped DP lookup still forwards correctly (DDP is only needed |
| 170 | + # for the backward all-reduce). Check the pooled (EBC) forward against a |
| 171 | + # local (unsharded) reference that mirrors the sharded module's actual |
| 172 | + # replicated weight -- the sharded module re-inits in reset_parameters, |
| 173 | + # so copy its state, not the source module's. The embedding kernels have |
| 174 | + # no deterministic CUDA variant, so relax the deterministic guard. |
| 175 | + if is_ebc: |
| 176 | + weight_key = f"embedding_bags.{_TABLE_NAME}.weight" |
| 177 | + local_module = _build_ebc(device) |
| 178 | + local_module.load_state_dict( |
| 179 | + {weight_key: frozen_dmp.state_dict()[weight_key].cpu()} |
| 180 | + ) |
| 181 | + kjt = _input_kjt(device) |
| 182 | + # Relax the deterministic guard only around the forward (embedding kernels |
| 183 | + # have no deterministic CPU/CUDA variant), then restore it -- the flag is |
| 184 | + # process-global even inside this MultiProcessContext subprocess. |
| 185 | + prev_deterministic = torch.are_deterministic_algorithms_enabled() |
| 186 | + torch.use_deterministic_algorithms(False) |
| 187 | + try: |
| 188 | + with torch.inference_mode(): |
| 189 | + frozen_out = _resolve(frozen_dmp(kjt)) |
| 190 | + local_out = local_module(kjt) |
| 191 | + torch.testing.assert_close( |
| 192 | + # pyre-ignore[16]: _resolve() returns object; the EBC forward output |
| 193 | + # (KeyedTensor) exposes .values(). |
| 194 | + frozen_out.values(), |
| 195 | + # pyre-ignore[16]: local EBC forward returns a KeyedTensor (.values()). |
| 196 | + local_out.values(), |
| 197 | + ) |
| 198 | + finally: |
| 199 | + torch.use_deterministic_algorithms(prev_deterministic) |
| 200 | + |
| 201 | + |
| 202 | +@skip_if_asan_class |
| 203 | +class DataParallelGradSyncOptOutTest(MultiProcessTestBase): |
| 204 | + def _run(self, is_ebc: bool) -> None: |
| 205 | + # Hide CUDA only for this test (not at module import, which would leak into |
| 206 | + # other test modules). The spawned workers inherit the env before they import |
| 207 | + # torch, forcing MultiProcessContext onto its CPU/gloo path. |
| 208 | + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): |
| 209 | + self._run_multi_process_test( |
| 210 | + callable=_run_dp_grad_sync_optout, |
| 211 | + world_size=2, |
| 212 | + backend="gloo", |
| 213 | + is_ebc=is_ebc, |
| 214 | + ) |
| 215 | + |
| 216 | + def test_ebc_dp_skip_grad_sync(self) -> None: |
| 217 | + self._run(is_ebc=True) |
| 218 | + |
| 219 | + def test_ec_dp_skip_grad_sync(self) -> None: |
| 220 | + self._run(is_ebc=False) |
| 221 | + |
| 222 | + |
| 223 | +class MarkDataParallelSkipGradSyncTest(unittest.TestCase): |
| 224 | + """Unit tests for the marker helpers (no sharding / distributed setup needed).""" |
| 225 | + |
| 226 | + def test_unmarked_module_returns_false(self) -> None: |
| 227 | + # A module never passed to mark_data_parallel_skip_grad_sync opts in to the |
| 228 | + # default (DDP reducer kept), so should_skip_data_parallel_grad_sync is False. |
| 229 | + module = _build_ebc(torch.device("cpu")) |
| 230 | + self.assertFalse(should_skip_data_parallel_grad_sync(module)) |
| 231 | + |
| 232 | + def test_marked_module_returns_true(self) -> None: |
| 233 | + module = _build_ebc(torch.device("cpu")) |
| 234 | + mark_data_parallel_skip_grad_sync(module) |
| 235 | + self.assertTrue(should_skip_data_parallel_grad_sync(module)) |
| 236 | + |
| 237 | + def test_unsharded_module_does_not_warn(self) -> None: |
| 238 | + module = _build_ebc(torch.device("cpu")) |
| 239 | + with self.assertNoLogs("torchrec.modules.embedding_modules", level="WARNING"): |
| 240 | + mark_data_parallel_skip_grad_sync(module) |
| 241 | + |
| 242 | + def test_already_sharded_module_warns(self) -> None: |
| 243 | + # Marking after sharding is a no-op: the marker is read once at shard time. |
| 244 | + # `_lookups` (set by ShardedEmbeddingModule) stands in for a sharded module |
| 245 | + # so this stays a unit test with no distributed setup. Assigned through |
| 246 | + # object.__setattr__ because nn.Module.__setattr__ is typed to accept only |
| 247 | + # Module | Tensor, while the real attribute is a plain List[nn.Module]. |
| 248 | + module = _build_ebc(torch.device("cpu")) |
| 249 | + object.__setattr__(module, "_lookups", []) |
| 250 | + with self.assertLogs( |
| 251 | + "torchrec.modules.embedding_modules", level="WARNING" |
| 252 | + ) as cm: |
| 253 | + mark_data_parallel_skip_grad_sync(module) |
| 254 | + self.assertIn("already-sharded", cm.output[0]) |
| 255 | + # Still sets the attribute -- the warning is advisory, not a guard. |
| 256 | + self.assertTrue(should_skip_data_parallel_grad_sync(module)) |
![meta-codesync[bot]](https://avatars.githubusercontent.com/in/1373969?v=4&size=40){kind=avatar}
0 commit comments