🚀[FEA]: Update FLARE++ to the v2 reference design and add token-sharded context parallelism for FLARE/FLARE++

Open
  • Is something wrong?
#2,024 2 comments 1 reaction 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
distributed-systems, machine-learning

Research direction

Start by reading the existing implementations in physicsnemo.nn.FLAREPlusPlus, physicsnemo.models.flare.FLAREPlusPlus, and GeoTransolver, then compare them with FLAREPPMixer in pdebench/models/mixer_backbone.py. For context parallelism, study pdebench/distributed/flare_cp.py and flare_encode in pdebench/models/flarepp.py alongside Appendices A and B of the linked paper. Done means implementing either the v2 mixer or token-sharded support for FLARE and FLAREPlusPlus without the current context-parallelism rejection.

Written by the indexing model from the issue text.

Description

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Medium

Please provide a clear description of problem you would like to solve.

cc @coreyjadams @mnabian

Following up on #1150 (FLARE), I wanted to share the updated version of FLARE++.

It's great to see FLARE++ already landed in #1940 (physicsnemo.nn.FLAREPlusPlus, physicsnemo.models.flare.FLAREPlusPlus, and GALE_FPP in GeoTransolver). That implementation follows the v1 preprint, so I want to clarify how v1 and v2 relate.

v1 vs. v2. v1 was preliminary work that we posted so there would be something public before my PhD defense. Because it was premature, we didn't release any code with it. v2 (our ICLR submission) is the polished version and has several improvements. What v1 calls FLARE++ is Simplified FLARE++ in v2 (Appendix A, "Simplified FLARE++"). The model v2 calls FLARE++ is the updated mixer described below.

Two things in v2 should be of interest to the team:

1. The FLARE++ token mixer as a replacement for FLARE

Reference: FLAREPPMixer in pdebench/models/mixer_backbone.py. Appendix A of v2 covers the design in full.

Changes from v1 (= the #1940 implementation) to v2:

  • Normalizing the synthesis call: head-wise normalization of the seeds and the synthesis keys/values (Appendix A, "Normalizing the synthesis call").
  • Fixed reference and gate: the routing query is a learned fixed query plus a per-head sigmoid-gated synthesized term, q = q_fixed + σ(g) · q_dynamic (Appendix A, "Fixed reference and gate"). With the gate closed, the mixer reduces to FLARE.
  • Near-identity feature projections: residual-MLP key/value projections initialized near identity (Appendix A, "Near-identity feature projections").
  • Three input projections instead of four: the synthesis keys and values are shared (Appendix A, "Sharing synthesis keys and values"). This saves one $C \to C$ projection per block without hurting accuracy in our experiments.

The normalization and initialization choices in Appendix A drastically improve early-training stability and help FLARE++ converge better. Appendix A also explains why synthesized routing queries are hard to optimize without them ("Optimization difficulty of synthesized routing queries"). "Operations in one block" in Appendix A lists the per-block operation sequence of FLARE, FLARE++, and Simplified FLARE++ side by side.

Two results from Appendix D support this:

  • Smaller seed variance (Appendix D.1): over ten seeds on Elasticity, FLARE++ has a smaller spread than FLARE. Its coefficient of variation is about 5% (0.37–0.43% test error), against about 11% for FLARE (0.58–0.86%).
  • Lower mixed-precision sensitivity (Appendix D.2): FLARE++ is the least precision-sensitive mixer in the sweep. Its FP16 error stays within 8% of its FP32 error at every depth.
2. Multi-GPU context parallelism for FLARE and FLARE++

Reference: pdebench/distributed/flare_cp.py. The model calls it through flare_encode in pdebench/models/flarepp.py. Appendix B of v2 has the complete derivation.

Point tokens are sharded across ranks, and no rank ever gathers the full sequence. The distributed encoder is exact: each rank runs a fused local SDPA against its shard. The rank then contributes only latent-sized outputs and log-sum-exp statistics to a collective, which rebuilds the globally normalized attention (FlareEncoderCPFlash / FlareFlashCPFunction, with a custom backward). Decoding stays local. FLARE calls this encoder once per block. FLARE++ calls it twice, once for query synthesis and once for the physical-value gather. Communication therefore depends on $M$ and the channel width, not on $N$.

Appendix B covers:

  • the token-sharded representation
  • the exact distributed encoder
  • the numerically stable reductions
  • the FLARE and FLARE++ communication schedules
  • the backward pass
  • measured scaling

#1940 added token-sharded FLARE++ and then reverted it, and it currently rejects context parallelism explicitly. This encoder could fill that gap for both FLARE and FLAREPlusPlus.

I'm happy to help with a PR for either piece, or to review.

Describe any alternatives you have considered

No response

Dominant language
Python
Stars
3.4k
Forks
828
Avg merge
4d 3h
Merged PRs (30d)
81

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/physicsnemo

All issues in NVIDIA/physicsnemo

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.