The trick is parameter groups!
optimizer = ZeroRedundancyOptimizer(
                rest_params,
                optim.AdamW,
                lr=args.lr,
                betas=(args.beta1, args.beta2),
                weight_decay=args.wd,
                eps=args.eps,
            )
optimizer.add_param_group({"params": gain_or_bias_params, "weight_decay": 0.})