pytorch - 💡(How to fix) Fix FSDP2 with ignore_params does not work with foreach optimizers [1 participants]

#!/usr/bin/env python3

import os

import torch
from torch import nn
from torch.distributed._composable.fsdp import fully_shard
from torch.distributed._tensor import DeviceMesh, DTensor


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.l = nn.Linear(4, 4)

    def forward(self, x):
        return self.l(x).sum()


def main():
    torch.distributed.init_process_group()
    rank = torch.distributed.get_rank()
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    use_cuda = torch.cuda.is_available()
    device = torch.device(f'cuda:{local_rank}' if use_cuda else 'cpu')
    if use_cuda:
        torch.cuda.set_device(device)

    model = Model().to(device)
    ignored = {model.l.bias}

    mesh = DeviceMesh.from_group(
        torch.distributed.GroupMember.WORLD,
        'cuda' if use_cuda else 'cpu',
    )
    fully_shard(model, mesh=mesh, ignored_params=ignored, reshard_after_forward=False)

    foreach = bool(int(os.environ.get('FOREACH', '1')))
    opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, foreach=foreach)
    x = torch.randn(2, 4, device=device)
    loss = model(x)
    loss.backward()

    if rank == 0:
        print(f'Calling optimizer.step() with foreach={foreach} ...')
    opt.step()
    if rank == 0:
        print('optimizer.step() succeeded')

    torch.distributed.destroy_process_group()


if __name__ == '__main__':
    main()

torchrun --standalone --nproc_per_node=2 repro_fsdp2_ignored_params_foreach.py