File size: 2,699 Bytes
f065e53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | def create_optimizer(model, weight_decay, learning_rate, betas=(0.9, 0.95)):
"""Create an AdamW optimizer with weight decay for 2D parameters only.
If the model defines `dit_lr_scale` (SpatialDiffuseSlot Phase-2 unfreeze),
the pretrained DiT trunk (dit.* except cond-embedder/null_cond) is placed in
separate param groups carrying `lr_scale` — timm's scheduler multiplies the
scheduled lr (warmup ramp AND cosine) by lr_scale per group, so the trunk
gets a gentler effective lr while new modules keep the full recipe.
"""
# start with all of the candidate parameters
param_dict = {pn: p for pn, p in model.named_parameters()}
# filter out those that do not require grad
param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
dit_scale = getattr(model, "dit_lr_scale", None)
def _is_trunk(n):
return (dit_scale is not None and n.startswith("dit.")
and not n.startswith(("dit.autoenc_cond_embedder", "dit.null_cond")))
# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2 and not _is_trunk(n)]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2 and not _is_trunk(n)]
optim_groups = [
{'params': decay_params, 'weight_decay': weight_decay},
{'params': nodecay_params, 'weight_decay': 0.0}
]
if dit_scale is not None:
trunk_decay = [p for n, p in param_dict.items() if p.dim() >= 2 and _is_trunk(n)]
trunk_nodecay = [p for n, p in param_dict.items() if p.dim() < 2 and _is_trunk(n)]
if trunk_decay or trunk_nodecay:
optim_groups += [
{'params': trunk_decay, 'weight_decay': weight_decay, 'lr_scale': dit_scale},
{'params': trunk_nodecay, 'weight_decay': 0.0, 'lr_scale': dit_scale},
]
if is_main_process():
n_t = sum(p.numel() for p in trunk_decay) + sum(p.numel() for p in trunk_nodecay)
print(f"[dit_lr_scale={dit_scale}] DiT trunk in lr-scaled groups: {n_t:,} params")
num_decay_params = sum(p.numel() for p in decay_params)
num_nodecay_params = sum(p.numel() for p in nodecay_params)
if is_main_process():
print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
optimizer = AdamW(optim_groups, lr=learning_rate, betas=betas)
return optimizer
|