Why you get 'expected all tensors to be on the same device', and where the mismatched tensor actually comes from

When code is first written for a single GPU, it is common to hardcode .to("cuda") or .cuda() on the model, inputs, or specific tensors. Once the model is later loaded with device_map='auto', which places different layers on different GPUs, that hardcoded call pins one tensor to a fixed device while the model expects inputs to match whatever device each layer actually landed on, causing a mismatch the moment that layer runs.

Fix: Remove hardcoded .to("cuda") calls on the model itself when using device_map='auto', and move inputs to the device of the model's embedding layer rather than a fixed index.

Collators often build auxiliary tensors from scratch, which default to CPU unless a device is explicitly specified. If the rest of the batch has already been moved to GPU by the training loop but the collator's output has not, the mismatch surfaces the first time the two tensors interact, such as in a loss or attention computation.

Fix: Either leave all collator output on CPU and let the Trainer or accelerate move the whole batch to the correct device afterward, or explicitly set the device to match the batch when the collator creates new tensors.

device_map='auto' uses the accelerate library to compute a per-layer device assignment based on available VRAM across GPUs, and expects nothing else to override that placement afterward. A subsequent .to("cuda:0") on the whole model attempts to move every parameter to a single device, which can partially succeed or silently break the mapping accelerate had already set up, leaving some tensors on their originally assigned device and others moved.

Fix: Never call .to(device) on a model loaded with device_map='auto'; if you need to move only the inputs, query the model's actual device placement instead of assuming cuda:0.

Some custom model implementations precompute and cache tensors like rotary embedding tables the first time they are needed, using whatever device is current at that moment. If that first call happens before multi-GPU placement is finalized, or from a different rank than later calls, the cached buffer stays pinned to the wrong device for subsequent calls.

Fix: Register such buffers with register_buffer() so they move automatically with the module's .to() calls, or explicitly move the cache to match the input tensor's device at each call rather than caching a device assumption.

Custom generation utilities sometimes create small tensors, such as a done flag or an EOS check tensor, which default to CPU, then combine them with GPU logits tensors during the generation loop, producing the same device mismatch pattern in a code path that is easy to overlook since it is not part of the main forward pass.

Fix: Explicitly pass a matching device to every tensor-creation call inside custom generation logic, matching whatever tensor it will be combined with.