[BUG]: FusedLAMB Fails with LowLevelZeroPlugin When Using Small initial_scale and CPU Offload

Is there an existing issue for this bug?

• I have searched the existing issues

The bug has not been fixed in the latest main branch

• I have checked the latest main branch

Do you feel comfortable sharing a concise (minimal) script that reproduces the error? :)

Yes, I will share a minimal reproducible script.

Describe the bug

Similar to https://github.com/hpcaitech/ColossalAI/issues/6400 ,When using LowLevelZeroPlugin, if initial_scale is set to a smaller value such as (2^{20}) (the default is (2^{32})) and cpu_offload=True, training with the FusedLAMB optimizer fails with the following error:RuntimeError: expected input to be on cuda

The specific reproduction script is main.py:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import colossalai
from colossalai.booster import Booster
from colossalai.nn.optimizer.fused_lamb import FusedLAMB
from colossalai.booster.plugin import LowLevelZeroPlugin

class RandomDataset(Dataset):
    def __init__(self, num_samples=32 * 10, input_dim=1024, num_classes=10):
        self.x = torch.randn(num_samples, input_dim)
        self.y = torch.randint(0, num_classes, (num_samples,))

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]


class MLP(nn.Module):
    def __init__(self, input_dim=1024, hidden_dim=512, num_layers=10, num_classes=10):
        super().__init__()
        layers = []
        for i in range(num_layers):
            in_dim = input_dim if i == 0 else hidden_dim
            layers.append(nn.Linear(in_dim, hidden_dim))
            layers.append(nn.ReLU())
        layers.append(nn.Linear(hidden_dim, num_classes))
        self.net = nn.Sequential(*layers)

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


def main():
    seed = 3377659546
    colossalai.launch_from_torch(seed=seed)
    plugin = LowLevelZeroPlugin(
        initial_scale=2**20,
        cpu_offload=True,
    )

    booster = Booster(plugin=plugin)

    model = MLP()
    optimizer = FusedLAMB(model.parameters(), lr=1e-3)
    print(optimizer)
    criterion = nn.CrossEntropyLoss()

    dataset = RandomDataset()
    train_dataloader = DataLoader(dataset, batch_size=32, shuffle=False)

    model, optimizer, criterion, train_dataloader, _ = booster.boost(model, optimizer, criterion, train_dataloader)

    precision = getattr(plugin, "precision", "fp16")
    dtype_map = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}
    dtype = dtype_map.get(precision, torch.float16)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.train()
    for epoch in range(1):
        total_loss = 0
        for step, (x, y) in enumerate(train_dataloader):
            x = x.to(device=device, dtype=dtype)
            y = y.to(device=device)

            optimizer.zero_grad()
            output = model(x)
            loss = criterion(output, y)

            booster.backward(loss, optimizer)
            optimizer.step()

            total_loss += loss.item()
            print(f"[Epoch {epoch}] step {step}, loss = {loss.item():.4f}")

        avg_loss = total_loss / len(train_dataloader)
        print(f"Epoch {epoch} finished, average loss = {avg_loss:.4f}")


if __name__ == "__main__":
    main()

Running the following training command:

colossalai run --nproc_per_node 4 --master_port 29505 ./main.py

Will produce the following error log:

[rank2]: Traceback (most recent call last):
[rank2]:   File "/home/yanzhen/distributed_test/colossalAI/test/./bug9.py", line 84, in <module>
[rank2]:     main()
[rank2]:   File "/home/yanzhen/distributed_test/colossalAI/test/./bug9.py", line 74, in main
[rank2]:     optimizer.step()
[rank2]:   File "/home/yanzhen/miniconda3/envs/colossal/lib/python3.9/site-packages/colossalai/zero/low_level/low_level_optim.py", line 588, in step
[rank2]:     self.optim.step()
[rank2]:   File "/home/yanzhen/miniconda3/envs/colossal/lib/python3.9/site-packages/torch/optim/optimizer.py", line 487, in wrapper
[rank2]:     out = func(*args, **kwargs)
[rank2]:   File "/home/yanzhen/miniconda3/envs/colossal/lib/python3.9/site-packages/colossalai/nn/optimizer/fused_lamb.py", line 133, in step
[rank2]:     g_norm_32 = multi_tensor_applier(self.multi_tensor_l2norm, self._dummy_overflow_buf, [g_all_32], False)[0]
[rank2]:   File "/home/yanzhen/miniconda3/envs/colossal/lib/python3.9/site-packages/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py", line 35, in __call__
[rank2]:     return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)
[rank2]: RuntimeError: expected input to be on cuda

Environment

PyTorch version: 2.5.1+cu124 Is debug build: False CUDA used to build PyTorch: 12.4 ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64) GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04.2) 11.4.0 Clang version: 18.1.3 (1ubuntu1) CMake version: version 3.28.3 Libc version: glibc-2.39

Python version: 3.9.23 (main, Jun 5 2025, 13:40:20) [GCC 11.2.0] (64-bit runtime) Python platform: Linux-6.5.0-18-generic-x86_64-with-glibc2.39 Is CUDA available: True CUDA runtime version: 12.4.99 CUDA_MODULE_LOADING set to: LAZY GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4090 GPU 1: NVIDIA GeForce RTX 4090 GPU 2: NVIDIA GeForce RTX 4090 GPU 3: NVIDIA GeForce RTX 4090

Nvidia driver version: 580.65.06 cuDNN version: Probably one of the following: /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn.so.8 /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8 /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_train.so.8 /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8 /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8 /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8 /usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_train.so.8 HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 48 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 256 On-line CPU(s) list: 0-255 Vendor ID: AuthenticAMD Model name: AMD EPYC 7773X 64-Core Processor CPU family: 25 Model: 1 Thread(s) per core: 2 Core(s) per socket: 64 Socket(s): 2 Stepping: 2 Frequency boost: enabled CPU max MHz: 3527.7339 CPU min MHz: 1500.0000 BogoMIPS: 4400.15 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm Virtualization: AMD-V L1d cache: 4 MiB (128 instances) L1i cache: 4 MiB (128 instances) L2 cache: 64 MiB (128 instances) L3 cache: 1.5 GiB (16 instances) NUMA node(s): 2 NUMA node0 CPU(s): 0-63,128-191 NUMA node1 CPU(s): 64-127,192-255 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Mitigation; safe RET Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Versions of relevant libraries: [pip3] galore-torch==1.0 [pip3] numpy==2.0.2 [pip3] torch==2.5.1 [pip3] triton==3.1.0 [conda] galore-torch 1.0 pypi_0 pypi [conda] numpy 2.0.2 pypi_0 pypi [conda] torch 2.5.1 pypi_0 pypi [conda] triton 3.1.0 pypi_0 pypi