trainer.test() with deepspeed bf16

[leao1995]

calling trainer.test() after trainer.fit() gives an error when using deepspeed and bf16.

ValueError: torch.float32 is enabled but the following parameters have dtype that is not torch.float32: [('module.model.weight', torch.bfloat16), ('module.model.bias', torch.bfloat16)]

deepspeed== 0.6.5
pytorch-lightning==1.6.4

The following is a minimul example to reproduce the error.

import torch
from torch.utils.data import Dataset
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.strategies.deepspeed import DeepSpeedStrategy
import logging

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.model = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        output = self.model(batch)
        loss = torch.sum(output)
        return loss

    def test_step(self, batch, batch_idx):
        return self.model(batch)

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]

train_data = torch.utils.data.DataLoader(RandomDataset(32, 64), num_workers=8, batch_size=8)
model = BoringModel()
strategy = DeepSpeedStrategy(
    config = {
        "bf16": { "enabled": True },
        "prescale_gradients": False,
        "zero_optimization": {
            "stage": 2,
            "contiguous_gradients": False,
            "allgather_bucket_size": 5e8,
            "reduce_bucket_size": 5e8,
            "overlap_comm": True
        },
        "zero_allow_untested_optimizer": True
    }
)
trainer = Trainer(
    limit_train_batches=1,
    limit_test_batches=1,
    max_epochs=1,
    devices=1,
    accelerator="gpu",
    precision="bf16",
    strategy=strategy
)
trainer.fit(model, train_data)
trainer.test(model, train_data)

[gg22mm]

我也遇到相同的问题，这个到底怎么解决 / I also encountered the same problem, how can I solve it？