Skip to content

Skip installing large aliased constant-folded weights - #4531

Open
micwill755 wants to merge 1 commit into
partition-short-circuitfrom
skip-aliased-fold-install
Open

Skip installing large aliased constant-folded weights#4531
micwill755 wants to merge 1 commit into
partition-short-circuitfrom
skip-aliased-fold-install

Conversation

@micwill755

Copy link
Copy Markdown
Collaborator

Summary

  • Skip installing large constant-folded tensors that only alias existing module weights, avoiding expensive cpu().contiguous() copies under offload_module_to_cpu.

PR 2 of 3 in a compile-startup stack; based on #4530.

Test plan

  • Flux NVFP4 cold compile still succeeds with offload_module_to_cpu=True
  • Large weight permutes remain in-graph (not installed as frozen params)
  • VLA still folds materialized embedding constants and compiles (int32 indices)

Avoid cpu().contiguous() copies of folded tensors that already share storage with module weights under offload_module_to_cpu, while still installing materialized folds.
@meta-cla meta-cla Bot added the cla signed label Aug 19, 2026
@github-actions github-actions Bot added component: lowering Issues re: The lowering / preprocessing passes component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 19, 2026
@github-actions
github-actions Bot requested a review from cehongwang August 19, 2026 21:56
@narendasan

Copy link
Copy Markdown
Collaborator

Can you explain more about what is happening here? What situation are we avoiding by skipping constant folding?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants