[BUG] Tensor are not on the same device when enable cpu activation offload

I am trying to enable cpu activation offload when training my custom LLAMA model. However, an error occur:

It seems like some inputs of the attention operation is offloaded in cpu but others are not. I thinks this is automatically handled by deepspeed, so do I use this feature in a wrong way?
My model code is like:

class Transformer(nn.Module):
    def __init__(self, params: ModelArgs, checkpoint_activations=False, checkpoint_num_layers=5):
        super().__init__()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers
        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)

        self.layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))

        self.norm = RMSNorm(params.dim, eps=params.norm_eps, name='last_norm')
        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)

        self.freqs_cis = precompute_freqs_cis(
            self.params.dim // self.params.n_heads, self.params.max_seq_len * 2
        )

        self.checkpoint_activations = checkpoint_activations
        self.checkpoint_num_layers = checkpoint_num_layers

        if deepspeed.checkpointing.is_configured():
            self.get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker
            self.checkpoint = deepspeed.checkpointing.checkpoint

    # @torch.inference_mode()
    def forward(self, tokens: torch.Tensor):

        def custom(start, end):
            def custom_forward(*inputs):
                layers_ = self.layers[start:end]
                x_ = inputs[0]
                for layer in layers_:
                    x_ = layer(x_, inputs[1], inputs[2])
                return x_

            return custom_forward

        _bsz, seqlen = tokens.shape
        start_pos = 0
        h = self.tok_embeddings(tokens)

        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos: start_pos + seqlen]

        mask = None
        if seqlen > 1:
            mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=tokens.device)
            mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)

        if self.checkpoint_activations:
            l = 0
            num_layers = len(self.layers)
            chunk_length = self.checkpoint_num_layers
            while l < num_layers:
                h = self.checkpoint(custom(l, l + chunk_length),
                               h, freqs_cis, mask)
                l += chunk_length
        else:
            for layer in self.layers:
                h, freqs_cis, mask = layer(h, freqs_cis, mask)
        h = self.norm(h)
        # (bsz, seq_len, vocab_size)
        output = self.output(h)
        return output.float()

And the deepspeed config is:

{
  "train_batch_size": 40,
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 10,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 3e-4,
      "eps": 1e-8,
      "weight_decay": 0.01

    }
  },
  "activation_checkpointing": {
    "partition_activations": true,
    "contiguous_memory_optimization": true,
    "cpu_checkpointing": true
  },
  "fp16": {
    "enabled": true,
    "min_loss_scale": 1,
    "loss_scale_window": 1000,
    "cpu_offload": true
  },
  "zero_optimization": {
    "stage": 3,
    "offload_param": {
    "device": "nvme",
    "nvme_path": "/foundation_model",
    "pin_memory": true
  },
    "offload_optimizer": {
    "device": "nvme",
    "nvme_path": "/foundation_model",
    "pin_memory": true
  },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": 26214400,
    "allgather_bucket_size": 1e8,
    "reduce_bucket_size": 1e8,
    "round_robin_gradients": false,
    "stage3_prefetch_bucket_size": 23592960,
    "stage3_param_persistence_threshold": 51200,
    "stage3_max_live_parameters": 2e8,
    "stage3_max_reuse_distance": 2e8,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "aio": {
      "block_size": 524288,
      "queue_depth": 8,
      "thread_count": 1,
      "single_submit": true,
      "overlap_events": true
    },

  "steps_per_print": 1
}

Hope someone can help me out, thank you!

The issue names no DeepSpeed source file or test; it provides a custom Transformer example and an activation_checkpointing configuration. Start by reproducing the device mismatch with that model and configuration, then trace the activation checkpointing and CPU offload paths. Done would require a confirmed cause plus either a targeted fix with regression coverage or clear configuration guidance.