device meta is invalid when using group offload

[bghira]

[RANK 0] 2025-10-21 13:15:19,710 [INFO] (simpletuner.helpers.models.common) Moving AutoencoderKLWan to accelerator, converting from torch.float32 to torch.bfloat16
[RANK 0] 2025-10-21 13:15:21,156 [ERROR] (validation) Error generating validation image: device meta is invalid, Traceback (most recent call last):
  File "/notebooks/SimpleTuner/simpletuner/helpers/training/validation.py", line 1675, in validate_prompt
    pipeline_result = self.model.pipeline(**pipeline_kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/simpletuner/helpers/models/wan/pipeline.py", line 1007, in __call__
    return self._call_text_to_video(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/simpletuner/helpers/models/wan/pipeline.py", line 643, in _call_text_to_video
    noise_pred_uncond = self.transformer(
                        ^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/diffusers/hooks/hooks.py", line 188, in new_forward
    args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/diffusers/hooks/group_offloading.py", line 304, in pre_forward
    self.group.onload_()
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/diffusers/hooks/group_offloading.py", line 260, in onload_
    self._onload_from_disk()
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/diffusers/hooks/group_offloading.py", line 189, in _onload_from_disk
    loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=device)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/notebooks/SimpleTuner/.venv/lib/python3.11/site-packages/safetensors/torch.py", line 381, in load_file
    with safe_open(filename, framework="pt", device=device) as f:
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
safetensors_rust.SafetensorError: device meta is invalid

encountered this bad boy when trying to use group offload with wan 2.1 14b t2v.

so my understanding of this, having hit it in HiDream before especially:

• the pipeline gets created with missing denoiser module (such as, for creating text embeds)
• the transformer is attached to the pipeline when it's loaded
• the pipeline still thinks device should be meta instead of the actual accelerator

and in hidream, 9d2cac2 fixed it by just overriding device = self.transformer.device

🤔 so maybe the same fix will work here. but it'd be nicer to understand the root issue. i've traced through it before, but not really looking forward to doing so again.

[bghira]

Loading text encoder and tokenizer...
Creating pipeline with transformer=None...
Attaching transformer to pipeline...
Pipeline _execution_device: cuda:0
Transformer device: cuda:0
Text encoder device: cuda:0

Pre-caching text encoder outputs...
Moving text encoder to 'meta' device...
Pipeline _execution_device after moving text encoder to meta: meta
Transformer device: cuda:0
Text encoder device: meta

Attempting to call the pipeline for validation...
  0%|                                                                                                                           | 0/2 [00:00<?, ?it/s]
Error occurred during pipeline call: Tensor on device cuda:0 is not on the expected device meta!
Error type: RuntimeError

import sys
import torch


def main():
    model_id = "pruna-test/wan-t2v-tiny-random"

    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("Loading text encoder and tokenizer...")
        from transformers import AutoTokenizer, UMT5EncoderModel

        tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer")
        text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", dtype=torch.bfloat16)

        from diffusers import AutoencoderKLWan, WanPipeline, FlowMatchEulerDiscreteScheduler

        vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)

        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")

        print("Creating pipeline with transformer=None...")
        pipeline = WanPipeline(
            tokenizer=tokenizer,
            text_encoder=text_encoder,
            transformer=None,  # Initially None
            vae=vae,
            scheduler=scheduler,
        )

        text_encoder.to(device)
        vae.to(device)

        from diffusers import WanTransformer3DModel

        transformer = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)
        transformer.to(device)

        print("Attaching transformer to pipeline...")
        pipeline.transformer = transformer

        print(f"Pipeline _execution_device: {pipeline._execution_device}")
        print(f"Transformer device: {transformer.device}")
        print(f"Text encoder device: {text_encoder.device}")

        print("\nPre-caching text encoder outputs...")
        prompt_embeds = pipeline.encode_prompt(
            prompt="A cat and a dog baking a cake together",
            device=device,
        )

        print("Moving text encoder to 'meta' device...")
        text_encoder.to("meta")

        # Check the pipeline's _execution_device after moving text encoder to meta
        print(f"Pipeline _execution_device after moving text encoder to meta: {pipeline._execution_device}")
        print(f"Transformer device: {transformer.device}")
        print(f"Text encoder device: {text_encoder.device}")

        try:
            print("\nAttempting to call the pipeline for validation...")
            try:
                result = pipeline(
                    # prompt="A cat and a dog baking a cake together",
                    height=320,
                    width=512,
                    num_frames=17,
                    num_inference_steps=2,  # Minimal steps for speed
                    guidance_scale=5.0,
                    prompt_embeds=prompt_embeds[0],
                    negative_prompt_embeds=prompt_embeds[1],
                )
                print("Pipeline call succeeded - no error reproduced.")
                return True
            except Exception as e:
                print(f"Error occurred during pipeline call: {e}")
                print(f"Error type: {type(e).__name__}")

        except Exception as e:
            print(f"Error occurred during group offloading: {e}")
            print(f"Error type: {type(e).__name__}")
            return False

    except ImportError as e:
        print(f"Import error: {e}")
        return False


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)

@sayakpaul here's a reproducer to attempt to provoke this issue, it gets closer but not to the exact bug. still, we easily reach the condition where execution_device is meta after moving the text encoder.

if i do this instead:

        print("Moving text encoder to 'meta' device...")
        pipeline.text_encoder = None
        text_encoder.to("meta")

we don't see the error anymore:

100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  4.35it/s]
Pipeline call succeeded - no error reproduced.

[bghira]

fixed with a workaround to patch the getter, but the bug persists upstream