improve parameter grouping/fusing logic

dinov2/train/ssl_meta_arch.py

@@ -368,24 +368,20 @@ def train(self):
         super().train()
         self.teacher.eval()
 
-    def get_maybe_fused_params_for_submodel(self, m):
-        params_groups = get_params_groups_with_decay(
-            model=m,
-            lr_decay_rate=self.cfg.optim.layerwise_decay,
-            patch_embed_lr_mult=self.cfg.optim.patch_embed_lr_mult,
-        )
-        fused_params_groups = fuse_params_groups(params_groups)
-        logger.info("fusing param groups")
-
-        for g in fused_params_groups:
-            g["foreach"] = True
-        return fused_params_groups
-
     def get_params_groups(self):
         all_params_groups = []
         for m in self.student.values():
-            all_params_groups += self.get_maybe_fused_params_for_submodel(m)
-        return all_params_groups
+            all_params_groups.extend(
+                get_params_groups_with_decay(
+                    model=m,
+                    lr_decay_rate=self.cfg.optim.layerwise_decay,
+                    patch_embed_lr_mult=self.cfg.optim.patch_embed_lr_mult,
+                )
+            )
+        fused_params_groups = fuse_params_groups(all_params_groups)
+        # sort param groups by their *sorted* param names to ensure deterministic ordering
+        fused_params_groups = sorted(fused_params_groups, key=lambda g: sorted(g["param_names"]))
+        return fused_params_groups
 
     def prepare_for_distributed_training(self):
         logger.info("DISTRIBUTED FSDP -- preparing model for distributed training")

dinov2/train/train.py

@@ -58,8 +58,8 @@ def get_args_parser(add_help: bool = True):
     return parser
 
 
-def build_optimizer(cfg, params_groups):
-    return torch.optim.AdamW(params_groups, betas=(cfg.optim.adamw_beta1, cfg.optim.adamw_beta2))
+def build_optimizer(cfg, params_groups, **kwargs):
+    return torch.optim.AdamW(params_groups, betas=(cfg.optim.adamw_beta1, cfg.optim.adamw_beta2), **kwargs)
 
 
 def build_schedulers(cfg):
@@ -138,7 +138,7 @@ def do_train(cfg, model, resume=False):
 
     # setup optimizer
 
-    optimizer = build_optimizer(cfg, model.get_params_groups())
+    optimizer = build_optimizer(cfg, model.get_params_groups(), foreach=True)
     (
         lr_schedule,
         wd_schedule,

dinov2/utils/param_groups.py

@@ -5,6 +5,7 @@
 
 from collections import defaultdict
 import logging
+from torch.distributed.fsdp._common_utils import clean_tensor_name
 
 
 logger = logging.getLogger("dinov2")
@@ -56,7 +57,7 @@ def get_params_groups_with_decay(model, lr_decay_rate=1.0, patch_embed_lr_mult=1
     all_param_groups = []
 
     for name, param in model.named_parameters():
-        name = name.replace("_fsdp_wrapped_module.", "")
+        name = clean_tensor_name(name)
         if not param.requires_grad:
             continue
         decay_rate = get_vit_lr_decay_rate(
@@ -80,14 +81,13 @@ def get_params_groups_with_decay(model, lr_decay_rate=1.0, patch_embed_lr_mult=1
 
 
 def fuse_params_groups(all_params_groups, keys=("lr_multiplier", "wd_multiplier", "is_last_layer")):
-    fused_params_groups = defaultdict(lambda: {"params": []})
+    fused_params_groups = defaultdict(lambda: {"params": [], "param_names": set()})
     for d in all_params_groups:
-        identifier = ""
+        identifier = "-".join([f"{k}={d[k]}" for k in keys])
+        group = fused_params_groups[identifier]
         for k in keys:
-            identifier += k + str(d[k]) + "_"
-
-        for k in keys:
-            fused_params_groups[identifier][k] = d[k]
-        fused_params_groups[identifier]["params"].append(d["params"])
-
-    return fused_params_groups.values()
+            group[k] = d[k]
+        group["params"].append(d["params"])
+        group["param_names"].add(d["name"])
+        group["name"] = identifier
+    return list(fused_params_groups.values())