support multiple outer optims for diloco

Summary:
- support passing in a different outer optimizer for each fragment
- currently accept both list of optimizers and a single optimizer for backward compatibility

Author: tushar00jain

torchft/local_sgd.py

@@ -550,7 +550,8 @@ def __init__(
         manager: Manager,
         model_fragments: List[nn.Module],
         inner_optimizer: optim.Optimizer,
-        outer_optimizer: optim.Optimizer,
+        # TODO: this is for backward compatibility
+        outer_optimizer: optim.Optimizer | list[optim.Optimizer],
         sync_every: int,
         backup_device: Optional[torch.device] = None,
         pin_memory: bool = True,
@@ -575,6 +576,11 @@ def __init__(
             fragment_update_alpha: Determines how to mix the local and global optimized parameters
         """
 
+        if isinstance(outer_optimizer, list):
+            assert len(outer_optimizer) == len(
+                model_fragments
+            ), "The number of outer optimizers must match the number of model fragments"
+
         if manager._use_async_quorum:
             raise ValueError(
                 "Using DiLoCo require synchronous quorum to be enabled. "
@@ -623,8 +629,11 @@ def __init__(
                 model_fragment,
                 math.floor((sync_every / len(model_fragments)) * (i + 1)),
                 inner_optimizer,
-                # TODO: Support different outer optimizers for each fragment
-                outer_optimizer,
+                (
+                    outer_optimizer[i]
+                    if isinstance(outer_optimizer, list)
+                    else outer_optimizer
+                ),
                 sync_every,
                 backup_device,
                 pin_memory,