DeepSpeed ​​Lightning отказывается распараллеливать слои даже при настройке на этап 3

Я хочу придумать очень простой пример Lightning с использованием DeepSpeed, но он отказался распараллеливать слои даже при настройке на этап 3.

Я просто расширяю модель, добавляя слои FC в надежде, что они будут распределены по разным графическим процессорам (всего 6)

Но я заканчиваю с

torch.cuda.OutOfMemoryError: CUDA не хватает памяти. Пытался выделить 2,00 МБ (GPU 3; общая емкость 15,00 ГБ; 14,00 ГБ уже выделено; 5,25 ГБ свободно; 14,00 ГБ зарезервировано в общей сложности PyTorch) Если зарезервированная память >> выделенная память, попробуйте установить max_split_size_mb во избежание фрагментации. См. документацию по управлению памятью и PYTORCH_CUDA_ALLOC_CONF

Поэтому я предполагаю, что слои помещаются только в один графический процессор.

Полный код доступен здесь, но вот суть:

Раздутие модели с 18000 слоями:

class TelModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        embed_dim = 512
        component_list = [
                nn.Linear(512, embed_dim)
        #] + [nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True) for _ in range(n_layers)] + [
        ] + [nn.Linear(embed_dim, 512) for _ in range(n_layers)] + [
                nn.Linear(embed_dim, 512)
        ]
        self.net = torch.nn.Sequential(*component_list)

Инициализация DeepSpeed:

tel_model = TelModel()
train_ds = RandomDataset(100)
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE)
trainer = L.Trainer(accelerator = "gpu", devices=6, strategy = "deepspeed_stage_3", precision=32)
trainer.fit(tel_model, train_loader)

И, наконец, я запускаю это так:

глубокая скорость молнии-deepspeed-tel.py