Make fused gate-up projection the default by anijain2305 · Pull Request #4535 · pytorch/torchtitan

anijain2305 · 2026-09-09T12:52:48Z

Stack from ghstack (oldest at bottom):

Human note

The main complicated change is the LoRA adaptors. The fused lora linear layer has the logical slices w1 and w3 still available, so that lora can be applied to each slice individually.

Agent summary

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion.

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: 9f06cad Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: da98279 Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: 73866be Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: 7c7ec0c Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: af364a8 Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: b53af9f Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: b69fd50 Pull-Request: #4535

[ghstack-poisoned]

Use one physical w13 linear for dense feed-forward gate and up projections while preserving logical w1 and w3 configs and checkpoint keys. Integrate the layout with DistGEMM, LoRA, quantization, Kimi, and serving consumers, and keep the Triton override focused on the activation fusion. ghstack-source-id: a2726ab Pull-Request: #4535

[ghstack-poisoned]

anijain2305 · 2026-09-09T23:00:07Z

-8 0.0002728958206716925
-9 -0.00023937312653288245
-10 -0.0019825994968414307
+2 0.002184871584177017


It changes because the backward accumulation order changes

Earlier, it was dX = dgrad_w1_path + dgrad_w3_path , and now its one big matrix multiplication.

tianyu-l · 2026-09-10T00:03:13Z

    return symm_mem


 class AllGatherLinear(torch.autograd.Function):


In @pianpwk's rewrite #4427, I think we need these concepts?

Maybe we should use these names for the "sync" version, and rename these modules to AsyncAllAgatherLinear, which is then applied as ModelConfigTransform, following #4322. @fegin

Let me take that up in a separate PR. I will learn about #4322.

tianyu-l · 2026-09-10T00:06:33Z

 class SigmoidGatedFeedForward(FeedForward):
    """SwiGLU feed-forward with a per-token sigmoid gate.

    The output is ``sigmoid(gate(x)) * ffn(x)``. Inherits ``w1/w2/w3`` from


no w1/w2/w3 any more

tianyu-l · 2026-09-10T00:09:30Z

    def __init__(self, config: Config):
        super().__init__()
-        self.w1 = config.w1.build()
+        self.w13 = _build_interleaved_linear(


looks a bit complicated -- why can't we mimic FusedQKV and make w13 a single Linear.Config?

tianyu-l · 2026-09-10T00:13:21Z

-                    layouts[f"{module_prefix}{state_name}"] = layout
-
-                # FusedSwiGLU exposes split w1/w3 state-dict keys while the
-                # layout is declared on the fused w13 parameter.
-                w13_layout = sharding_config.state_shardings.get("w13")
-                if w13_layout is not None:
-                    for proj_name in ("w1", "w3"):
-                        layouts[f"{module_prefix}{proj_name}.weight"] = w13_layout


btw this is not ideal, but the change doesn't look any better haha.

Also we are also doing the same for FusedQKVLinear. Ideally we should let each fused module own such conversions instead of hardcoding them in this function.

Actually, I start to doubt the value of these state dict hooks -- in the past, it was because we host both fused and unfused impls and we want the checkpoint to be interchangeable. But now that we are removing unfused impls, there's no need to stick with unfused state dict? After all it will cause inefficiency like #4549 (cc @tushar00jain)

I wonder if we could just

remove all hooks

adapt all the state_dict_adapters.py to handle the HF format conversion "natively"

no need to worry about fused grouped experts

This indeed sounds cleaner. Can I take this up in a separate PR?

tianyu-l · 2026-09-10T00:19:35Z

For LoRA, I somehow is not convinced that we have to make w1 w3 separate.

If I think about why users would use LoRA instead of original weights

the main benefit is on memory reduction, and also some throughput gain

they don't really care 100% of the numerics with original impl

Both of them are still true if we just do LoRA on the fused w13, LoL?

I think we should just do lora on fused w13 and get rid of all this complexity. I will do some more search and take that path if nothing serious comes up

tianyu-l · 2026-09-10T00:26:27Z

-            _situ_glu(self.w1(x), self.w3(x), self.beta, self.linear_beta),
-        )
+    def _activation(self, gate_TF: torch.Tensor, up_TF: torch.Tensor) -> torch.Tensor:
+        return _situ_glu(gate_TF, up_TF, self.beta, self.linear_beta)


instead of this, I think a better encapsulation for training framework is to make activation_fn: ActivationFn a configurable function owning beta and linear_beta, subclassing https://github.com/pytorch/torchtitan/blob/main/torchtitan/config/function.py#L16

This way we don't need to define KimiFeedForward and KimiGroupedExperts at all. And we can put all activation functions in common/activation.py

[ghstack-poisoned]

Update

9f5631b

[ghstack-poisoned]

pytorch-bot Bot added the ciflow/rl label Sep 9, 2026

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

Update

2c59dc7

[ghstack-poisoned]

Update

95a9f77

[ghstack-poisoned]

Update

feeed6b

[ghstack-poisoned]

Update

fe0bce5

[ghstack-poisoned]

Update

447b377

[ghstack-poisoned]

Update

28bd0f2

[ghstack-poisoned]

Update

d181e05

[ghstack-poisoned]

Update

ebb2d90

[ghstack-poisoned]

anijain2305 marked this pull request as ready for review September 9, 2026 17:22

anijain2305 requested review from fegin, tianyu-l, wconstab and wwwjn as code owners September 9, 2026 17:22

This was referenced Sep 9, 2026

Make fused QKV projection the default implementation #4546

Open

Add RegionAC feed-forward regions #4526

Open

[DONT REVIEW] Add RegionAC grouped-expert regions #4550

Draft

Make fused grouped-expert projection the default #4552

Open

Update

ab3a396

[ghstack-poisoned]

anijain2305 commented Sep 9, 2026

View reviewed changes

tianyu-l reviewed Sep 10, 2026

View reviewed changes

anijain2305 added 3 commits September 9, 2026 18:09

Update

36dafac

[ghstack-poisoned]

Update

c0235fc

[ghstack-poisoned]

Update

12a60f7

[ghstack-poisoned]

Navigation Menu

Uh oh!

Make fused gate-up projection the default - #4535

Make fused gate-up projection the default#4535
anijain2305 wants to merge 13 commits into
gh/anijain2305/12/basefrom
gh/anijain2305/12/head

anijain2305 commented Sep 9, 2026 •

edited

Loading

Uh oh!

anijain2305 Sep 9, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

anijain2305 Sep 10, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

anijain2305 Sep 10, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

anijain2305 Sep 10, 2026

Uh oh!

tianyu-l Sep 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return symm_mem


		class AllGatherLinear(torch.autograd.Function):

Uh oh!

Conversation

anijain2305 commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Human note

Agent summary

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

Labels

Projects

Milestone

Development

Uh oh!

2 participants

anijain2305 commented Sep 9, 2026 •

edited

Loading