set flag instead

Author: casteryh

src/forge/actors/generator.py

@@ -272,12 +272,10 @@ def split_keys(keys):
         futures = []
         for i, names in enumerate(split_keys(hf_param_names)):
 
-            async def fetch_coro():
-                return self.weight_fetchers.slice(procs=i).fetch.call_one(
-                    version=version, param_names=names
-                )
+            fut = self.weight_fetchers.slice(procs=i).fetch.call_one(
+                version=version, param_names=names
+            )
 
-            fut = asyncio.create_task(fetch_coro())
             futures.append(fut)
 
         sub_state_dicts = [await fut for fut in futures]

src/forge/env.py

@@ -105,6 +105,17 @@ def get_value(self) -> Any:
     description="Whether or not to use RDMA in TorchStore.",
 )
 
+MONARCH_OLD_ASYNC_WORKAROUND = EnvVar(
+    name="MONARCH_OLD_ASYNC_WORKAROUND",
+    default=1,
+    description=(
+        "If enabled, monarch messages will be sent immediately even it's not"
+        " awaited. This is needed for parallel fetching of weights, as using"
+        " create_task creates race condition. This is a temporary workaround"
+        " and will be removed once we have a better solution."
+    ),
+)
+
 
 def all_env_vars() -> list[EnvVar]:
     """Retrieves all registered environment variable names."""