mem_tokens receive gradients but are excluded from the optimizer in DDP training

[1FirstWave1]

Description

In meta_train_parallel.py, mem_tokens are marked as trainable and receive non-zero gradients during training, but they are excluded from all optimizer parameter groups when DDP is enabled.

As a result:

• mem_tokens.grad is computed;
• the gradient norm keeps increasing;
• optimizer.step() never updates mem_tokens;
• optimizer.zero_grad(set_to_none=True) does not clear their gradients;
• the values of mem_tokens remain unchanged throughout training.

This appears inconsistent with the SHINE paper, which lists the initial memory embeddings as trainable parameters.

Relevant code

utils/myfreeze.py explicitly enables gradients for mem_tokens:

def freeze(metamodel):
    for param in metamodel.parameters():
        param.requires_grad = False
    metamodel.model.mem_tokens.requires_grad = True

However, the optimizer parameter groups in meta_train_parallel.py exclude every parameter whose name starts with module.metamodel:

    {
        "params": [
            p for n, p in ddp_metanet.named_parameters()
            if (
                not any(nd in n for nd in no_decay)
                and not n.startswith("module.metamodel")
            )
        ],
        "weight_decay": cfg.optim.weight_decay,
    },
    {
        "params": [
            p for n, p in ddp_metanet.named_parameters()
            if (
                any(nd in n for nd in no_decay)
                and not n.startswith("module.metamodel")
            )
        ],
        "weight_decay": 0.0,
    },
    {
        "params": list(
            iter_learnable_tensors(metalora)
            if not USE_ADDITIONAL_METALORA
            else iter_learnable_tensors(ift_additional_metalora)
        ),
        "weight_decay": cfg.optim.weight_decay,
    },
]

Under DDP, the full parameter name of the memory embeddings is:

module.metamodel.model.mem_tokens

Reproduction / observed behavior

I added the following check after optimizer initialization:

raw_metanet = (
    ddp_metanet.module
    if isinstance(ddp_metanet, DDP)
    else ddp_metanet
)
mem_tokens = raw_metanet.metamodel.model.mem_tokens

mem_in_optimizer = any(
    mem_tokens is param
    for group in optimizer.param_groups
    for param in group["params"]
)

print(
    f"requires_grad={mem_tokens.requires_grad}, "
    f"in_optimizer={mem_in_optimizer}"
)

The result is:

requires_grad=True, in_optimizer=False

I also compared the values immediately before and after optimizer.step() and monitored the gradient norm.
The observed behavior is:

grad_norm: increasing and non-zero
delta_norm: 0
delta_max: 0
changed_elements: 0

This shows that backward propagation reaches mem_tokens, but the optimizer never updates them.
Because they are not part of any optimizer parameter group, the following call also does not clear their gradients:

optimizer.zero_grad(set_to_none=True)

This explains why their gradient norm continues to increase.

[Yewei-Liu]

Yes, we actually noticed this bug before as well. Thanks a lot for digging so deeply into the code!

We also conducted some ablation studies and found that freezing the initial embedding of the memory token does not have a significant impact on the results. In fact, keeping it frozen seems to make the training slightly more stable. So leaving it frozen is totally fine.

So it's nice to see this issue and you are right :)