Regional weather diffusion: torch.compile with domain parallelism

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

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)
 

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,

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,

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)