Loading DeepSeek R1 model took extremely long time

[Neo9061]

System Info

Following the recent merged PR and release note, I try to load the DeepSeek R1 model using code snippet below on single P5EN (8 H200) GPUs.

1. The first issue I have is that it took very long time to load. Estimation is ~10 hours.
2. Then I tried to modify the config.json and model.safetensors.index.json to try to load first 10 layers including embed_token and lm.head modules. However, it hit following error. The issue is gone if I used DeepSeek conversion script to convert the checkpoint from FP8 to BF16.

Some parameters are on the meta device because they were offloaded to the cpu.
The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:1 for open-end generation.
The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Traceback (most recent call last):
  File "/iofsx/sds3/models/DeepSeekV3/test.py", line 18, in <module>
    outputs = model.generate(inputs, max_new_tokens=50)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/generation/utils.py", line 2370, in generate
    result = self._sample(
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/generation/utils.py", line 3331, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/models/deepseek_v3/modeling_deepseek_v3.py", line 1025, in forward
    outputs = self.model(
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/models/deepseek_v3/modeling_deepseek_v3.py", line 773, in forward
    layer_outputs = decoder_layer(
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/models/deepseek_v3/modeling_deepseek_v3.py", line 513, in forward
    hidden_states, self_attn_weights = self.self_attn(
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/transformers/models/deepseek_v3/modeling_deepseek_v3.py", line 423, in forward
    q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(query_shape).transpose(1, 2)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/opt/conda/envs/fix/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::BFloat16 != c10::Float8_e4m3fn

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Code snippet to load and generate text example.

# `run_deepseek_v1.py`
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
torch.manual_seed(30)
model_path = "MYMODEL_PATH"
tokenizer = AutoTokenizer.from_pretrained(model_path)

chat = [
  {"role": "user", "content": "Hello, how are you?"},
  {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
  {"role": "user", "content": "I'd like to show off how chat templating works!"},
]


model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.bfloat16)
inputs = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)

outputs = model.generate(inputs, max_new_tokens=50)
print(tokenizer.batch_decode(outputs))

Expected behavior

NA

[ArthurZucker]

cc @Cyrilvallez can you have a look?

[spensun-aws]

Hi, I monitored another smiliar issue. To bypass the issue mentioned here, I used the torchrun and made it to load the model fast, however, the generation is extremely slow. It takes around 20 min to generate 500 tokens on a 8xH200. When I attempt to generate a long response from sharegpt dataset samples with 4096 max tokens, it hangs forever, seems would take much time to generate.

cmd I used: torchrun --nproc_per_node=8 gen.py

gen.py content:

#!/usr/bin/env python3
import json
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Set a manual seed for reproducibility.
torch.manual_seed(30)

def main():
    inputfile="xx.jsonl"
    output="xxx.jsonl"
    model_path="/xxx/model/DeepSeek-R1"

    # Load tokenizer and model.
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, device_map="auto", torch_dtype=torch.bfloat16
    )

    with open(inputfile, 'r') as infile, open(output, 'w') as outfile:
        for line in infile:
            line = line.strip()
            if not line:
                continue

            try:
                data = json.loads(line)
            except json.JSONDecodeError as e:
                print(f"Skipping line due to JSON error: {e}")
                continue

            # Ensure the 'conversation' key exists and is a list.
            chat = data["conversations"]
            if not isinstance(chat, list):
                print("Skipping line: 'conversation' is not a list.")
                continue

            # Prepare the inputs using the conversation.
            inputs = tokenizer.apply_chat_template(
                chat,
                tokenize=True,
                add_generation_prompt=True,
                return_tensors="pt"
            ).to(model.device)

            # Generate model output.
            outputs = model.generate(inputs, max_new_tokens=4096,do_sample=False)
            # Decode the generated text.
            decoded_output = tokenizer.batch_decode(outputs)[0]

            # Append the generated output to the conversation.
            chat.append({"gpt": decoded_output})
            data["conversation"] = chat

            # Write the updated JSON object to the output file.
            outfile.write(json.dumps(data) + "\n")

if __name__ == '__main__':
    main()

[bruno686]

I also have this problem. It's extremely slow.

[ArthurZucker]

This is somehwat expected super happyu to see contributions! !!!

[maximizemaxwell]

Does this problem still occurs? I would like to try on it

[ArthurZucker]

Should be fixed if you correctly use torchrun !