RuntimeError: Input type (float) and bias type (c10::Half) should be the same

[PABannier]

System Info

Ubuntu 22.04.5 LTS

Python 3.10

dependencies = [
    "accelerate>=1.2.0",
    "deepspeed>=0.16.1",
    "numpy>=2.2.0",
    "setuptools>=75.6.0",
    "torch>=2.5.1",
    "transformers>=4.47.0",
]

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

1. Copy this dummy training script:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from accelerate import Accelerator
from pathlib import Path


# Dummy dataset that generates random data
class DummyDataset(Dataset):
    def __init__(self, size=1000):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        # Generate random image (3 channels, 64x64) and random label
        image = torch.randn(3, 64, 64).float()
        label = torch.randint(0, 2, (1,))[0]
        return image, label

# Simple CNN model
class DummyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 16 * 16, 64),
            nn.ReLU(),
            nn.Linear(64, 2)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

def main():
    # Initialize accelerator
    accelerator = Accelerator()

    # Create model and move to device
    model = DummyModel()

    # Create datasets and dataloaders
    train_dataset = DummyDataset(size=1000)
    train_loader = DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=True,
        num_workers=2
    )

    # Setup training components
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # Prepare for distributed training
    model, optimizer, train_loader = accelerator.prepare(
        model, optimizer, train_loader
    )

    # Training loop
    num_epochs = 5

    for epoch in range(num_epochs):
        model.train()
        total_loss = 0

        for batch_idx, (images, labels) in enumerate(train_loader):
            optimizer.zero_grad()

            outputs = model(images)
            loss = criterion(outputs, labels)

            accelerator.backward(loss)
            optimizer.step()

            total_loss += loss.item()

            if batch_idx % 10 == 0 and accelerator.is_main_process:
                print(f'Epoch: {epoch+1}/{num_epochs}, Batch: {batch_idx}, Loss: {loss.item():.4f}')

        if accelerator.is_main_process:
            avg_loss = total_loss / len(train_loader)
            print(f'Epoch: {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}')

if __name__ == "__main__":
    main()

2. Copy this DeepSpeed Zero3 configuration file:

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 12,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "500000000",
        "stage3_prefetch_bucket_size": "50000000",
        "stage3_param_persistence_threshold": "100000",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

3. Copy this Accelerate config file:

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  deepspeed_config_file: zero3_offload_config_accelerate.json
  zero3_init_flag: false
distributed_type: DEEPSPEED
downcast_bf16: "no"
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

4. Run:

uv run accelerate launch --config_file default.yaml ./dummy_train.py

I end up with this error: [rank0]: RuntimeError: Input type (float) and bias type (c10::Half) should be the same. I can fix it by explicitly casting the input tensor with half() but Accelerate/DeepSpeed should deal with mixed precision training, right?

Am I doing something wrong?

Expected behavior

Accelerate automatically handles mixed precision training, and correctly casts the input and weights to the correct type.