[New Model]Add DeepSeek V4 model support by Matrix-Z97 · Pull Request #3634 · pytorch/torchtitan

Matrix-Z97 · 2026-06-11T12:48:12Z

Description

This PR adds deepseek_v4 model support and registers it in TorchTitan’s supported model list.

Main Changes

Add the new torchtitan/models/deepseek_v4/ model directory
Implement the DeepSeek V4 model architecture, sparse attention, compressor, MoE, mHC, sharding, parallelization, and state dict adapter
Add a deepseek_v4 debug model configuration
Register deepseek_v4 in torchtitan/models/init.py

Parallelism Support

FSDP is currently supported
EP (Expert Parallel) is currently supported
TP (Tensor Parallel) is currently supported
CP (Context Parallel) is not supported yet and explicitly raises NotImplementedError

Additional Notes

The attention implementation is currently based on small operators rather than a fused/custom kernel.
Only the debug model configuration is included for now.

If needed, I can further improve and complete this implementation.

pytorch-bot · 2026-06-11T12:48:19Z

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

Matrix-Z97 · 2026-06-12T08:52:34Z

@tianyu-l I noticed that you added deepseek_v4 to the roadmap. However, I'm sorry that I just force-pushed a new version of the code, which caused the roadmap content to disappear. Could you please add it back?

tianyu-l · 2026-06-12T08:56:50Z

@Matrix-Z97 you mean the "26H1 TorchTitan Development"?

Oh that's not a roadmap but just for tracking what this repo is doing. I re-marked the status from "Done" to "In Progress".

Matrix-Z97 · 2026-06-12T09:03:32Z

@tianyu-l Understood, thanks for the correction! Out of curiosity, are you interested in the deepseek_v4 model?

tianyu-l · 2026-06-12T09:10:09Z

Yes, of course. We have plans, but were lacking bandwidth.

This PR comes in very timely, and the quality looks very high. We just need some time to review. Many thanks!

Btw @shuhuayu is working on adding latest Kimi model.

Matrix-Z97 · 2026-06-15T08:46:34Z

Thanks a lot for the feedback! Happy to help with the bandwidth. Let me know if there's anything I can do to help move this PR forward or if you need any adjustments. And awesome news about the latest Kimi model @shuhuayu.

tianyu-l · 2026-06-17T00:24:35Z

cc @drisspg to review the attention part

Co-authored-by: Jonathan Zhou <96200235+floatingtrees@users.noreply.github.com>

Adds model/loss-agnostic machinery for model-internal auxiliary losses: - AuxLossInjection / inject_aux_loss — identity-forward autograd.Function that injects an auxiliary loss gradient into a carrier tensor without changing the forward value (PP-safe, fullgraph-compile-safe, single-tensor output contract). SPMD-typed via @spmd.register_autograd_function. - reduce_to_replicate — backend-agnostic differentiable Partial->Replicate reduction over named mesh axes (cp/tp). Dispatches on tensor type: DTensor -> full_dtensor redistribute; plain -> spmd_types redistribute. Single-process / no-mesh calls are pure-local no-ops. - ModelSpec.metrics_fn — generic optional callback for model-internal metrics collection, called by the trainer during metrics collection. Default None (no behavior change for existing models).

floatingtrees · 2026-06-29T20:11:04Z

@Matrix-Z97 Are you planning to add MTP support for this model? If not, I’d be happy to add it.

Matrix-Z97 · 2026-07-02T02:15:48Z

@floatingtrees Thanks, MTP has definitely been on our roadmap, and we’ve also been iterating on it continuously in our internal versions. Our current thinking is more along the lines of providing DeepSeek-V3 support in the spirit of #3392. Do you have any better ideas or suggestions?

tianyu-l

having some high level feedback before looking into details

tianyu-l · 2026-07-07T05:46:36Z

+        ({parallel_dims.tp}) and 2 * CP degree ({parallel_dims.cp}).
+        """
+
+    if parallelism.spmd_backend in ("full_dtensor", "spmd_types"):


We are in the migration to "spmd_types". Could you remove "default" and "full_dtensor" support from this new model?

Understood. So this new model only needs to support the spmd backend for now? Does this mean all torchtitan models will transition to supporting only the spmd backend in the future?

tianyu-l · 2026-07-07T05:48:04Z

+deepseek_v4_configs = {
+    "debugmodel": _debugmodel,
+}


why only debug model, can we include the full set?

I've included deepseek_v4_flash and deepseek_v4_pro now to expand the set. Just a heads-up: due to hardware constraints, I haven't actually tested them yet.

tianyu-l · 2026-07-07T05:50:13Z

    pipelining_fn: Callable | None
    post_optimizer_build_fn: Callable | None
    state_dict_adapter: type[BaseStateDictAdapter] | None
+    metrics_fn: Callable | None = None


@felipemello1 do you have a plan on in-model metric logging that we could share?

Let's add a TODO in dsv4 model code, and remove the model logging in this PR.

tianyu-l · 2026-07-07T05:56:16Z

+
+
+@spmd.register_autograd_function
+class AuxLossInjection(torch.autograd.Function):


I see similar contribution in #3864

Can we make sure the design in #3000 suits the need for both?

Yes, this implementation was referenced from #3864. I'll need to look closer into #3000 to see if the design can cover both use cases. Will update this thread once I have more clarity.

From a quick look, #3000 seems more like a MoE-specific interface, while here we would like to call the aux-loss capability through a common/shared interface that can also cover non-MoE losses such as the DSA indexer loss.

So for the current DeepSeek V4 use case, the #3864-style implementation seems more suitable to me.

I'll take a look at both

tianyu-l · 2026-07-07T05:58:04Z

could you please add verification (e.g. at least in PR summary) following https://github.com/pytorch/torchtitan/tree/main/scripts/checkpoint_conversion#comprehensive-check-kl-divergence

cc @shuhuayu

Thanks for the suggestion. Running the KL check with deepseek_v4_flash or deepseek_v4_pro as the baseline is not feasible for me due to hardware limits.

Would it be acceptable to run the same HF-vs-TorchTitan KL / max-diff comparison on a custom small DeepSeek V4 config instead?

Got it. @shuhuayu since you've done it for Kimi, could you help verify this PR gets to similar level of numerical parity?

Would it be acceptable to run the same HF-vs-TorchTitan KL / max-diff comparison on a custom small DeepSeek V4 config instead?

We also did a customized 16b kimi configs to compare numerics so it sounds good to me to do the same for dsv4. @tianyu-l, do you think we need to test it in the original 862b configs, doable but needs some additional setups.

Thanks. Could you share what config was used for the customized 16B Kimi numerical test? In particular, which dimensions were reduced, e.g. number of layers, hidden size, heads/head dim, MoE experts, or sequence length?
I can follow the same strategy for DSV4 and build a smaller custom config for the KL/numerical comparison.

I see. Small scale verification sounds fine.

Thanks. Could you share what config was used for the customized 16B Kimi numerical test?

I used the kimi-vl variant, see https://github.com/pytorch/torchtitan/pull/3532/changes#diff-5a13a40a31333dd8b095a114b15d9580c2744ac09211c6b551a02c2989032c9b

I used the kimi-vl variant, see https://github.com/pytorch/torchtitan/pull/3532/changes#diff-5a13a40a31333dd8b095a114b15d9580c2744ac09211c6b551a02c2989032c9b
Got it, will try it out soon. Thanks!

Matrix-Z97 · 2026-07-10T07:46:27Z

        return torch.cat((-x2, x1), dim=-1)


+class SingleComplexRoPE(ComplexRoPE):


@tianyu-l Hi, it looks like the existing community RoPE interface only supports (query, key) pair inputs, while some DeepSeek V4 RoPE call sites operate on a single tensor, such as in the compressor and indexer. Without a single-tensor API, these paths would need awkward patterns like rope(x, x)[0].

To avoid that, we added SingleComplexRoPE for explicit single-tensor complex RoPE usage. It also supports inverse=True, which is needed to match the HF inference logic in the DeepSeek V4 attention post phase.

Does this form of change look acceptable to the community?

@shuhuayu I vaguely remember there's a model using single-field rope, but I couldn't find it any more. My worry is that after recent refactor the single-field rope is gone. Was it originally added by you or do you have clue?

I searched and did not find a single field one neither. rather than a new SingleComplexRoPE class, how about just make key optional in the existing RoPE.forward / apply_rotary_emb (return a single tensor when key is None) and add an inverse=False flag ? That covers both the single-tensor call sites and the inverse-rope post-phase without a parallel class, sharing one implementation, and a no-op for existing models.

@shuhuayu
my only concern is that it also need to work with the helion kernels (https://github.com/pytorch/torchtitan/blob/main/torchtitan/overrides/helion_rope.py), but maybe AI can just do it

This is not urgent for dsv4 since it uses ComplexRoPE and helion kernels is only for CosSineRoPE now. but it leaves the work to people who wants to add helion ComplexRoPE kernels.

@shuhuayu
See #3767

I think at least we shouldn't break existing key, query pair version of Helion override.

Matrix-Z97 requested review from fegin, tianyu-l, wconstab and wwwjn as code owners June 11, 2026 12:48

pytorch-bot Bot added the ciflow/8gpu label Jun 11, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 11, 2026

tianyu-l added this to 26H1 TorchTitan Development Jun 11, 2026

tianyu-l moved this to In Progress in 26H1 TorchTitan Development Jun 11, 2026

Matrix-Z97 closed this Jun 12, 2026

Matrix-Z97 force-pushed the deepseek_v4 branch from fda1fcc to 07828a4 Compare June 12, 2026 08:29

github-project-automation Bot moved this from In Progress to Done in 26H1 TorchTitan Development Jun 12, 2026

Matrix-Z97 reopened this Jun 12, 2026

Add deepseek_v4 model support

90c57d1

Matrix-Z97 force-pushed the deepseek_v4 branch from daf4565 to 90c57d1 Compare June 12, 2026 08:39

tianyu-l moved this from Done to In Progress in 26H1 TorchTitan Development Jun 12, 2026

tianyu-l reviewed Jun 17, 2026

View reviewed changes

Comment thread torchtitan/models/deepseek_v4/attention.py

tianyu-l Jun 17, 2026

Copy link
Copy Markdown

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @drisspg to review the attention part

tianyu-l mentioned this pull request Jun 26, 2026

[WIP] DeepSeek-V3.2 model support #3800

Draft

1 task

floatingtrees reviewed Jun 26, 2026

View reviewed changes

Comment thread torchtitan/models/deepseek_v4/compressor.py Outdated

Matrix-Z97 and others added 4 commits June 27, 2026 09:37

Update torchtitan/models/deepseek_v4/compressor.py

c2d8a7a

Co-authored-by: Jonathan Zhou <96200235+floatingtrees@users.noreply.github.com>

Merge branch 'pytorch:main' into deepseek_v4

9a47905

Refactor DeepSeek V4 sparse attention routing

2991579

Add DeepSeek V4 indexer aux loss

7811710

tianyu-l reviewed Jul 7, 2026

View reviewed changes

Update DeepSeek V4 configs and sharding

811953e

Matrix-Z97 requested review from IvanKobzarev, SherlockNoMad, aditvenk, felipemello1, sanketpurandare and xmfan as code owners July 10, 2026 07:23

Refactor DeepSeek V4 attention, RoPE, and aux loss

7d49617

Matrix-Z97 force-pushed the deepseek_v4 branch from 0627776 to 7d49617 Compare July 10, 2026 07:34

Matrix-Z97 commented Jul 10, 2026

View reviewed changes



		@spmd.register_autograd_function
		class AuxLossInjection(torch.autograd.Function):

		return torch.cat((-x2, x1), dim=-1)


		class SingleComplexRoPE(ComplexRoPE):

Uh oh!

Conversation

Matrix-Z97 commented Jun 11, 2026

Description

Main Changes

Parallelism Support

Additional Notes

Uh oh!

pytorch-bot Bot commented Jun 11, 2026

Uh oh!

Matrix-Z97 commented Jun 12, 2026

Uh oh!

tianyu-l commented Jun 12, 2026

Uh oh!

Matrix-Z97 commented Jun 12, 2026

Uh oh!

tianyu-l commented Jun 12, 2026

Uh oh!

Matrix-Z97 commented Jun 15, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

floatingtrees commented Jun 29, 2026

Uh oh!

Matrix-Z97 commented Jul 2, 2026

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees