diff --git a/CHANGELOG.md b/CHANGELOG.md index 163df3ccfe..b09f703d67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -129,6 +129,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- The regional weather diffusion recipe compiles its loss with + `torch.compile` under domain parallelism too (it was turned off), and takes + `attention_backend` from the model config. - Partitioning a graph across GPUs (`partition_graph_with_id_mapping`, behind `DistributedGraph` and GraphCast's distributed mode) no longer loops over destination nodes in Python. A 1-million-node grid now takes well under a diff --git a/examples/weather/regional_weather_diffusion/test_training.py b/examples/weather/regional_weather_diffusion/test_training.py index c699a1bd5e..2dabb166a3 100644 --- a/examples/weather/regional_weather_diffusion/test_training.py +++ b/examples/weather/regional_weather_diffusion/test_training.py @@ -120,8 +120,7 @@ def test_distribute_model_syncs_before_sharding(monkeypatch): "domain_parallel_size, force_sharding", [(1, False), (1, True), (2, False)] ) @pytest.mark.parametrize("fp_optimizations", ["fp32", "amp-bf16"]) -# @pytest.mark.parametrize("torch_compile", [True, False]) -@pytest.mark.parametrize("torch_compile", [False]) +@pytest.mark.parametrize("torch_compile", [True, False]) @pytest.mark.parametrize("scheduler", [None, "CosineAnnealingLR"]) @pytest.mark.parametrize("sigma_distribution", ["lognormal", "loguniform"]) def test_training( @@ -153,12 +152,6 @@ def test_training( pytest.skip( f"Skipping: not enough processes ({dist.world_size}) to use domain_parallel_size of {domain_parallel_size}." ) - sharding = (domain_parallel_size > 1) or force_sharding - if sharding and torch_compile: - pytest.skip( - "Skipping: torch.compile is not supported with ShardTensor for now." - ) - # Set up rundir in the temporary directory rundir = _setup_rundir(tmp_path, dist.world_size) diff --git a/examples/weather/regional_weather_diffusion/utils/nn.py b/examples/weather/regional_weather_diffusion/utils/nn.py index 06986ae7b4..542dca6f3a 100644 --- a/examples/weather/regional_weather_diffusion/utils/nn.py +++ b/examples/weather/regional_weather_diffusion/utils/nn.py @@ -147,7 +147,7 @@ def get_preconditioned_natten_dit( depth=depth, num_heads=num_heads, patch_size=patch_size, - attention_backend="natten2d", + attention_backend=model_kwargs.pop("attention_backend", "natten2d"), layernorm_backend=layernorm_backend, attn_kwargs=attn_kw, condition_dim=condition_dim, diff --git a/examples/weather/regional_weather_diffusion/utils/parallel.py b/examples/weather/regional_weather_diffusion/utils/parallel.py index f71fd7e682..a99c7b697e 100644 --- a/examples/weather/regional_weather_diffusion/utils/parallel.py +++ b/examples/weather/regional_weather_diffusion/utils/parallel.py @@ -138,7 +138,12 @@ def __init__( mesh_shape=(self.data_parallel_size, domain_parallel_size), mesh_dim_names=["ddp", "domain"], ) - self.domain_rank = self.mesh["domain"].get_local_rank() + # Slice each submesh once and reuse the object: every mesh["..."] builds a + # new, equal DeviceMesh, and a compiled graph that meets one it does not + # track fails in Inductor (pytorch/pytorch#201701). + self.domain_mesh = self.mesh["domain"] + self.ddp_mesh = self.mesh["ddp"] + self.domain_rank = self.domain_mesh.get_local_rank() self.use_shard_tensor = use_shard_tensor def get_domain_group_zero_rank(self) -> int: @@ -149,7 +154,7 @@ def get_domain_group_zero_rank(self) -> int: int Global rank for local domain rank 0. """ - return torch.distributed.get_global_rank(self.mesh["domain"].get_group(), 0) + return torch.distributed.get_global_rank(self.domain_mesh.get_group(), 0) def local_batch_size(self, global_batch_size: int) -> int: """Compute per-rank batch size for data parallelism. @@ -339,7 +344,7 @@ def sharded_data_iter( i = 0 batch = None - domain_group = self.mesh["domain"].get_group() + domain_group = self.domain_mesh.get_group() while True: # the source rank within the domain group (always 0 when domain_parallel_size == 1) source_rank_in_mesh = i % self.domain_parallel_size @@ -431,10 +436,10 @@ def distribute_model(self, model: torch.nn.Module) -> FSDPModule: if not p.is_contiguous(): p.data = p.data.contiguous() - sync_module_over_mesh(model, self.mesh["domain"]) + sync_module_over_mesh(model, self.domain_mesh) if self.use_shard_tensor: self._shard_spatial_params(model) - fully_shard(model, mesh=self.mesh["ddp"]) + fully_shard(model, mesh=self.ddp_mesh) return model def _shard_spatial_params(self, model: torch.nn.Module) -> None: @@ -451,7 +456,7 @@ def _shard_spatial_params(self, model: torch.nn.Module) -> None: model : torch.nn.Module Model whose spatial parameters/buffers should be sharded in place. """ - domain_mesh = self.mesh["domain"] + domain_mesh = self.domain_mesh for module in model.modules(): for name, param in list(module.named_parameters(recurse=False)): shard_dim = shard_dim_selector(name) @@ -500,7 +505,7 @@ def make_domain_parallel_scheduler(self, scheduler: object) -> object: return DomainParallelNoiseScheduler( scheduler, - self.mesh["domain"], + self.domain_mesh, shard_dim=self.shard_dim, ) @@ -523,7 +528,7 @@ def replicate_tensor(self, t: torch.Tensor) -> torch.Tensor: if not self.use_shard_tensor or isinstance(t, DTensor): return t return DTensor.from_local( - t, device_mesh=self.mesh["domain"], placements=[Replicate()] + t, device_mesh=self.domain_mesh, placements=[Replicate()] ) def sync_replicated_grads(self, model: torch.nn.Module) -> None: @@ -553,7 +558,7 @@ def sync_replicated_grads(self, model: torch.nn.Module) -> None: """ if not self.use_shard_tensor: return - mesh = self.mesh["domain"] + mesh = self.domain_mesh mesh_size = mesh.size(0) if mesh_size == 1: return @@ -623,7 +628,7 @@ def nested_scatter( x = scatter_tensor( x, global_rank_of_source, - self.mesh["domain"], + self.domain_mesh, placements=(placement,), # Shard along height (H dimension) global_shape=x.shape, dtype=x.dtype, diff --git a/examples/weather/regional_weather_diffusion/utils/trainer.py b/examples/weather/regional_weather_diffusion/utils/trainer.py index a48e2575a2..2c51c676e7 100644 --- a/examples/weather/regional_weather_diffusion/utils/trainer.py +++ b/examples/weather/regional_weather_diffusion/utils/trainer.py @@ -683,12 +683,7 @@ def _setup_loss(self) -> WeightedMSEDSMLoss | RegressionLoss: self.sampling_scheduler = None - compile_loss = self.use_torch_compile and not self.use_shard_tensor - if self.use_torch_compile and self.use_shard_tensor: - self.logger.info( - "Skipping torch.compile on loss: not supported with " - "domain parallelism / ShardTensor." - ) + compile_loss = self.use_torch_compile if self.loss_type == "regression": loss_fn = RegressionLoss(self.net)