[asynctp] Knobs to enable asynctp; Adding constraints to solver to

incentivize asynctp

Author: IvanKobzarev

autoparallel/api.py

@@ -160,6 +160,7 @@ def __init__(
         enable_ac: bool = True,
         # None means 'auto'
         ac_stage_size_in_GiB: Optional[Union[float, str]] = "auto",
+        enable_asynctp: bool = False,
         **kwargs,
     ):
         self.stack = ExitStack()
@@ -190,6 +191,8 @@ def __init__(
         self.enable_ac = enable_ac
         self.ac_stage_size_in_GiB = ac_stage_size_in_GiB
 
+        self.enable_asynctp = enable_asynctp
+
         # NB: rest of the construction happens in __enter__
         self.active = False
 
@@ -216,6 +219,7 @@ def __enter__(self):
             self.mesh,
             rescale_grad_comm_cost_for_mp,
             repeated_subgraphs=self.kwargs.get("repeated_subgraphs", False),
+            enable_asynctp=self.enable_asynctp,
         )
 
         # makes sharding of params and gradients the same