code/trainer_utils_create_optimizer.py

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