torchtitan icon indicating copy to clipboard operation
torchtitan copied to clipboard
verified publisher icon pytorch

[RFC][WIP][CP] Enable FlexAttention CP for llama3

Open fegin opened this issue 2 months ago • 0 comments

Stack from ghstack (oldest at bottom):

  • -> #1857
  • #1939

This PR uses the latest CP APIs to enable FlexAttention + CP for llama3. This PR removes the usage of context_paralle() context manager and use _context_parallel_shard() to shard the input data.

fegin avatar Oct 12 '25 05:10 fegin