DSV4: Add shard_indexer_acts to shard the CSA indexer scores across heads - #4909
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the shard_indexer_acts configuration flag to enable sharding of DeepSeek-V4 CSA indexer score activations over the logical axis. The implementation updates the attention layers to conditionally apply logical sharding to activations like query, compressed key-value, scores, and index scores during non-autoregressive modes, and includes a new unit test suite to verify this behavior. The review feedback identifies several clean-up opportunities in the new test file, specifically pointing out an unused import (logical_to_mesh_axes), an unused constant (SCORE_AXES), and an unused helper function (custom_rule_sets) along with its associated imports.
| import yaml | ||
| from flax import nnx | ||
| from flax.linen import partitioning as nn_partitioning | ||
| from flax.linen.partitioning import logical_to_mesh_axes |
There was a problem hiding this comment.
The import logical_to_mesh_axes is unused in this test file and can be safely removed.
| SEQ = 64 | ||
| RATE = 4 | ||
| HEADS = 4 | ||
| SCORE_AXES = ("activation_batch", "activation_heads", "activation_length", None) |
There was a problem hiding this comment.
The constant SCORE_AXES is defined but never used in this test file and can be safely removed.
| def custom_rule_sets(): | ||
| """Loads logical-axis rules from each custom mesh preset.""" | ||
| out = {} | ||
| for path in sorted(CUSTOM_RULE_DIR.glob("*.yml")): | ||
| with path.open("r", encoding="utf-8") as file: | ||
| rules = yaml.safe_load(file)["logical_axis_rules"] | ||
| out[path.stem] = [(name, tuple(axes) if isinstance(axes, list) else axes) for name, axes in rules] | ||
| return out |
There was a problem hiding this comment.
The function custom_rule_sets is defined but never used in the test suite. To keep the codebase clean and maintainable, please remove this unused definition along with its associated unused imports (yaml on line 24, pathlib.Path on line 17, and CUSTOM_RULE_DIR on line 45).
systalyze-ai
changed the title
…ss heads Constrains the indexer's quadratic per-head score tensor over the activation_heads axis. Selections and masks are unchanged. Without it, seq 16384 exceeds the v6e per-chip memory limit by 4.42 GiB. Co-authored-by: Sudarsanan <[email protected]> Co-authored-by: Armin <[email protected]> Co-authored-by: utlz <[email protected]>
47be565 to
65a91ba
Compare
2 participants
Description
Adds
shard_indexer_actsto shard the DeepSeek-V4 CSA indexer's quadratic score tensor acrossactivation_heads. This reduces peak memory at long context.Performance
Disabling
shard_indexer_actsat seq 16384 exceeds the per-chip memory limit by 4.42 GiB (35.67 GiB needed, 31.25 GiB available) and the run does not compile; with it enabled the run fits. Setup:ici_expert_parallelism=16,ici_tensor_parallelism=8,ici_fsdp_parallelism=1,max_target_length=16384,per_device_batch_size=0.125(global batch 16),LIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, dynamic-splash attention path enabled.Tests
tests/unit/indexer_activation_sharding_test.py— 9 passed.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.