Reproduce/enable DeepSeek R1 Distill Llama 8B

[mergennachin]

🚀 The feature, motivation and pitch

This task is to enable DeepSeek R1 Distill Llama 8B on ExecuTorch. That way, people can run these models in a mobile app, locally, without talking to the server.

In theory, ExecuTorch already supports Llama 3.1 8B architecture anyway, so it should just work out of the box (https://github.com/pytorch/executorch/blob/main/examples/models/llama/README.md)

Please document (and make necessary changes) on how to run DeepSeek R1 Distill Llama 8B e2e via ExecuTorch on iOS and Android.

Update 1:

Was able to verify that export works as such: #7981 (comment)

Update 2:

Currently looking into tokenizers

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @cccclai @helunwencser @dvorjackz @byjlw

[jackzhxng]

This should just work out of the box if we convert the distilled weights and use them in export_llama

[mergennachin]

This should just work out of the box if we convert the distilled weights and use them in export_llama

Yeah, we have a script to do that

https://github.com/pytorch/executorch/blob/main/examples/models/llama/UTILS.md#download-models-from-hugging-face-and-convert-from-safetensor-format-to-state-dict

[raziel]

Hi @mergennachin , would love to see a script exporting the deep seek model so we can give it a try with our delegates as well :-)

[CypherpunkSamurai]

I have a distilled model but I need help converting this to pte file from pth.

The config file required a little cleaning or the export_llama kept complaining about ModelArgs.__init__() got an unexpected keyword argument 'architectures'

Now the script has mismatch for layers layers.0.attention.wk.weight: copying param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]). error.

The shapes of all layers are different.

Any ideas how I can fix this?

[iseeyuan]

@CypherpunkSamurai It may be related to the compressed projection used in MLA. I'm creating the interface in #8039 so that different attentions can be implemented and added.

[CypherpunkSamurai]

@iseeyuan thank you very much for your work man. Let me know if I can help contribute.

Here's the error logs btw:

# removing these json keys fixes the ModelArgs issue
unsupported_args = [
    "architectures",
    "attention_bias",
    "attention_dropout",
    "bos_token_id",
    "eos_token_id",
    "hidden_act",
    "hidden_size",
    "initializer_range",
    "intermediate_size",
    "max_position_embeddings",
    "mlp_bias",
    "model_type",
    "num_attention_heads",
    "num_hidden_layers",
    "num_key_value_heads",
    "pretraining_tp",
    "rms_norm_eps",
    "rope_scaling",
    "tie_word_embeddings",
    "torch_dtype",
    "transformers_version",
    "use_cache"
]

Error logs:

!cd executorch \
    && python -m examples.models.llama.export_llama \
        --checkpoint "/tmp/deepseek-r1-llama-8b-pth/deepseek-r1-llama-8b.pth" \
        --params "/tmp/deepseek-r1-llama-8b-pth/config_filtered.json" \
        --output_name "deepseek-r1-llama-8b.pte" \
        -kv \
        --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32

Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/kaggle/working/executorch/examples/models/llama/export_llama.py", line 32, in <module>
    main()  # pragma: no cover
  File "/kaggle/working/executorch/examples/models/llama/export_llama.py", line 28, in main
    export_llama(args)
  File "/kaggle/working/executorch/examples/models/llama/export_llama_lib.py", line 533, in export_llama
    builder = _export_llama(args)
  File "/kaggle/working/executorch/examples/models/llama/export_llama_lib.py", line 668, in _export_llama
    builder_exported = _prepare_for_llama_export(args).export()
  File "/kaggle/working/executorch/examples/models/llama/export_llama_lib.py", line 565, in _prepare_for_llama_export
    _load_llama_model(
  File "/kaggle/working/executorch/examples/models/llama/export_llama_lib.py", line 940, in _load_llama_model
    EagerModelFactory.create_model(
  File "/kaggle/working/executorch/examples/models/model_factory.py", line 44, in create_model
    model = model_class(**kwargs)
  File "/kaggle/working/executorch/examples/models/llama/model.py", line 235, in __init__
    missing, unexpected = self.model_.load_state_dict(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 2581, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for Transformer:
	size mismatch for layers.0.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.0.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.0.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.0.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.0.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.1.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.1.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.1.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.1.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.1.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.2.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.2.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.2.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.2.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.2.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.3.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.3.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.3.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.3.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.3.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.4.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.4.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.4.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.4.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.4.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.5.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.5.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.5.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.5.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.5.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.6.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.6.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.6.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.6.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.6.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.7.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.7.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.7.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.7.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.7.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.8.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.8.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.8.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.8.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.8.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.9.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.9.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.9.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.9.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.9.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.10.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.10.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.10.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.10.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.10.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.11.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.11.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.11.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.11.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.11.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.12.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.12.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.12.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.12.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.12.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.13.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.13.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.13.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.13.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.13.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.14.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.14.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.14.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.14.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.14.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.15.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.15.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.15.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.15.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.15.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.16.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.16.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.16.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.16.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.16.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.17.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.17.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.17.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.17.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.17.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.18.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.18.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.18.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.18.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.18.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.19.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.19.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.19.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.19.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.19.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.20.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.20.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.20.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.20.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.20.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.21.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.21.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.21.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.21.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.21.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.22.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.22.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.22.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.22.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.22.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.23.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.23.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.23.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.23.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.23.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.24.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.24.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.24.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.24.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.24.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.25.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.25.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.25.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.25.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.25.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.26.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.26.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.26.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.26.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.26.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.27.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.27.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.27.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.27.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.27.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.28.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.28.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.28.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.28.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.28.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.29.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.29.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.29.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.29.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.29.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.30.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.30.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.30.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.30.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.30.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.31.attention.wk.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.31.attention.wv.weight: copying a param with shape torch.Size([1024, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096]).
	size mismatch for layers.31.feed_forward.w1.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).
	size mismatch for layers.31.feed_forward.w2.weight: copying a param with shape torch.Size([4096, 14336]) from checkpoint, the shape in current model is torch.Size([4096, 11008]).
	size mismatch for layers.31.feed_forward.w3.weight: copying a param with shape torch.Size([14336, 4096]) from checkpoint, the shape in current model is torch.Size([11008, 4096]).

[mergennachin]

@CypherpunkSamurai

can you try using this file params.json instead of config_filtered.json in your export script?

https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/blob/main/original/params.json

[CypherpunkSamurai]

@mergennachin same error. Layer sizes are not matching.

[mergennachin]

Hmm interesting, @CypherpunkSamurai (and @raziel fyi) here's how I was able to export and create a pte file just now

Step 1: Follow this step to set up ExecuTorch: https://pytorch.org/executorch/main/getting-started-setup
Step 2: Run examples/models/llama/install_requirements.sh for Llama specific requirements.
Step 3: Download the model

pip install -U "huggingface_hub[cli]"
huggingface-cli download deepseek-ai/DeepSeek-R1-Distill-Llama-8B --local-dir ~/models/deepseek-ai/DeepSeek-R1-Distill-Llama-8B --local-dir-use-symlinks False

Step 4: Convert the model to pth file.

pip install torchtune

and run this command

from torchtune.models import convert_weights
from torchtune.training import FullModelHFCheckpointer
import torch

# Convert from safetensors to TorchTune. Suppose the model has been downloaded from Hugging Face
checkpointer = FullModelHFCheckpointer(
    checkpoint_dir='/Users/mnachin/models/deepseek-ai/DeepSeek-R1-Distill-Llama-8B',
    checkpoint_files=['model-00001-of-000002.safetensors', 'model-00002-of-000002.safetensors'],
    output_dir='/tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/' ,
    model_type='LLAMA3' # or other types that TorchTune supports
)

print("loading checkpoint")
sd = checkpointer.load_checkpoint()

# Convert from TorchTune to Meta (PyTorch native)
sd = convert_weights.tune_to_meta(sd['model'])

print("saving checkpoint")
torch.save(sd, "/tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/checkpoint.pth")

Step 5: Download this file https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/blob/main/original/params.json and save into /tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/params.json

Step 6: Export to .pte file

python -m examples.models.llama.export_llama \
    --checkpoint /tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/checkpoint.pth \
	-p /tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/params.json \
	-kv \
	--use_sdpa_with_kv_cache \
	-X \
	-qmode 8da4w \
	--group_size 128 \
	-d fp16 \
	--metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}' \
	--embedding-quantize 4,32 \
	--output_name="DeepSeek-R1-Distill-Llama-8B.pte"

You can continue with the runtime and iOS/Android integration here: https://github.com/pytorch/executorch/blob/main/examples/models/llama/README.md

[CypherpunkSamurai]

@mergennachin I'm using all nightly versions of torchtune and torchao, and executorch as well. I'm unsure what's causing the issue but sure I'll try this method and update you tomorrow, it's like a few hours past midnight here 🥲

Thank you for the quick reply btw 😄👍🏻

[CypherpunkSamurai]

@mergennachin the script examples.models.llama.export_llama uses around 29 GB of RAM and 0 VRAM for some reason. I'm using kaggle and it causes the session to interrupt when the usage is out of limits.

Is this normal? Is there any way to limit memory usage when loading and converting the model?

I'm using the same code:

!cd /tmp/executorch && \
    python -m examples.models.llama.export_llama \
    --checkpoint "/tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/checkpoint.pth" \
	-p "/tmp/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/original/params.json" \
	-kv \
	--use_sdpa_with_kv_cache \
	-X \
	-qmode 8da4w \
	--group_size 128 \
	-d fp32 \
	--metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}' \
	--embedding-quantize 4,32 \
	--output_name="/kaggle/working/DeepSeek-R1-Distill-Llama-8B.pte"