Workaround

extensions/thunder/strategies/thunder_ddp.py

@@ -136,8 +136,13 @@ def setup_module(self, module: Module) -> Module:
                     "You already called `thunder.jit()` and generated an execution trace. It's too late to apply the"
                     " DDP transform. Remove the `forward` call before `fabric.setup()`"
                 )
-            # modify the reference
-            cd.fn = thunder.distributed.ddp(cd.fn, **self._ddp_kwargs)
+            assert cd.is_module  # sanity check
+            ddp_module = thunder.distributed.ddp(cd.fn, **self._ddp_kwargs)
+            # update the compile data state
+            cd.fn = ddp_module
+            assert hasattr(cd, "_processed_function")  # sanity check
+            cd._processed_function = ddp_module
+            cd.process_group_for_ddp = ddp_module.process_group_for_ddp
             return module
         else:
             module = thunder.distributed.ddp(module, **self._ddp_kwargs)

extensions/thunder/strategies/thunder_fsdp.py

@@ -167,14 +167,19 @@ def setup_module(self, module: Module) -> Module:
                     "You already called `thunder.jit()` and generated an execution trace. It's too late to apply the"
                     " FSDP transform. Remove the `forward` call before `fabric.setup()`"
                 )
-            # modify the reference
-            cd.fn = thunder.distributed.fsdp(
+            assert cd.is_module  # sanity check
+            fsdp_module = thunder.distributed.fsdp(
                 cd.fn,
                 device=self.root_device,
                 sharding_strategy=self.sharding_strategy,
                 bucketing_strategy=self.bucketing_strategy,
                 **self._fsdp_kwargs,
             )
+            # update the compile data state
+            cd.fn = fsdp_module
+            assert hasattr(cd, "_processed_function")  # sanity check
+            cd._processed_function = fsdp_module
+            cd.process_group_for_ddp = fsdp_module.process_group_for_ddp
             return module
         else:
             module = thunder.distributed.fsdp(