Support explicit max sequence lengths in attention wrappers - #62
Closed
Li Dong (donglixp) wants to merge 2063 commits into
Closed
Li Dong (donglixp) wants to merge 2063 commits into
Li Dong (donglixp) wants to merge 2063 commits into
Conversation
1. Add a flag to IRTensor to indicate whether it is originally a scalar tensor. 2. During graph transformation, do as it is. 3. When generate code, check the flag to generate correct code. unit test pass parity check pass
MiniTrainer workable version. the parity check is against lightning.
parity matched between lightning version & mini-trainer version
add mixed precision f16 optimizer
Loss is a special tensor in the computation graph. - requires_grad = True - the forward graph and backward graph share exactly a same tensor physically The main branch exists problem when partitioning the loss. Since the loss is a scalar tensor by default, it is partitioned along the value dimension. Assume we have a operator `nll_loss([1024, 2048], [1024]) -> [1]` with annotation `N+ C^, C^ ->1`. In LLM training, `N` is the token dim, `C` is the dictionary dim, partition along `N` will partition the loss along value. In the main branch, following code will be generated  Although it is runnable and correct, it breaks our definition of `IRSegment`, **the intermediate variable `nll_loss_10138` should not be passed out as an output tensor**. However, removing this sub-tensor directly does not solve the problem, since the real loss tensor is generated by an adapter `nnscaler.runtime.adapter.all_reduce`, which means its `requires_grad` field equals to `False` at runtime. In addition, the additional partitioned `nll_loss_10138` disappears at pipeline in the main branch.  Root causes are - when `gen_activations` is called to generate adapters, the returned adapter for the partitioned loss is wrong. It should be a `nnscaler.runtime.adapter.nn.allreduce_identity` instead of `nnscaler.runtime.adapter.all_reduce` - an additional compiling pass `Grouping` is called for spmd/tp. `Grouping` will dispatch the partitioned graph to each device and build an `IRSegment` for each device. - in the `create_segment` method, there is an additional check when determining the outputs: `isinstance(otensor, IRSubTensor) and otensor.is_loss()`. This check will add both of `nll_loss_10138` and `nll_loss_1955` to the segment's output. - `nll_loss_1955` is annotated with `requires_grad=False` and `grad=None`, `nll_loss_10138` is annotated `requires_grad=True` and `grad = gtensorxxx`. According to the logic in `get_backward_callsite_io_tensors`, `nll_loss_10138` will be recognized as the real loss to the backward graph. - However, in the pipeline code generation, there is no `Grouping` pass. The dispatch process (ExeReuseCell -> Segment -> IRCell) strictly follows the assumption that output of a segment should be a full tensor. To solve this problem, in this PR - generate correct adapters when the output loss is used in another operator (like the `.data` operation in fairseq's criterion) - choose tensor as the segment's output carefully to make the emit process runnable parity check passed 
Azure pipeline agent is limited to 10GB space, and cannot be changed. The current solution is to delete virtual env after every test config is done
…it code scalar tensor: use right shape in generated module init code unit test pass parity check pass
…not be saved) fix checkpoint bug
…odist interface - add an example code to reduce the memory footprint when the sequence length and dictionary size is extremely large. It is verified in real model training - add an option `transient_mem_coef` to control the memory constraint
Verified on llama3 8B + 4K on 4xA6000, distributed plan is 2 pipeline stages, each stage is composed of 2 devices.  This PR includes: 1. refine autodist implementations, including - add option `parallel_profile` to control whether profiling nodes in parallel: in pipeline solver we need to build the SPMDSolver and profile nodes for many times, only the first constructing needs parallel profiling to speed up, dumping the graph and sync takes a lot of time when in parallel -> we only profile in serial for later SPMDSolver constructing - fix bugs to generate correct partition plans and analysis for intervals whose searching result is built from identical intervals - add a flag in gecode to tell front-end the loaded module is a pipeline stage or not. It is helpful when compile stage is separated from runtime. 2. fix bug in executor to support bf16 backward 3. refine comments unit test passed parity test passed with [PR](https://dev.azure.com/msrasrg/SuperScaler/_git/Fairseq/pullrequest/2220) verified on 8xH100 with 4 pipeline stages, each stage is composed of one stage.
minitrainer: add document
…global compute config parallel module: remove pipeline related config from global compute config
cross_entropy supports `reduction='none'` which means do no reduction. The PR adds support for this case. https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html
…pport. minitrainer: Add model/optimizer/lr_scheduler hook support.
1. register pre-hook to each reducer to divide scaling_factor 2. adjust back grads after allreduce
in this pr, directly copy torch 2.3 pytree as `_pytree.py`, it is more powerful than the previous version, and implemented most of the functions we needed, the only reason to copy the file instead of import is we also support version 2.0 <= torch < 2.3. there is no need to review `_pytree.py`. the additional support for pytree is in `pytree_utils.py` parity check passed.
Minitrainer logging: log tag support
1. Wrap batchnorm2d/instancenorm2d as a customized function because it has control flow in its forward function. Create new batchnorm2d/instancenorm2 module for replacing the original modules using a utility function automatically. 2. support communication within operator. An example operator is batchnorm2d when partitioning the batch dimension.
minitrainer: remove torchrun requirements for compile
… debug and optimize for autodist output each operator's importance ratio (percentages of states that can be reduced by forcing the operator to be partitioned in a single partition) an example output is ```text operator FwOp7-()(name=embedding, inputs=(t1768(p20,(1, 8192),d(),v(0/1)), w1770(p22,(32256, 4096),d(),v(0/1))), outputs=(t1771(p24,(1, 8192, 4096),d(),v(0/1)),)) has 4 partitions, importance ratio 0.225 at File "/home/yizhu1/ts_dev/Fairseq/nnscaler_examples/finetune_hf_model/src/model_helper/customize/modeling_nnscaler_mixtral_4_42.py", line 1047, in forward, inputs_embeds = self.embed_tokens(input_ids) operator FwOp22-()(name=transpose, inputs=(t1786(p66,(1, 8192, 8, 128),d(),v(0/1)),), outputs=(t1787(p68,(1, 8, 8192, 128),d(),v(0/1)),)) has 3 partitions, importance ratio 0.159 at File "/home/yizhu1/ts_dev/Fairseq/nnscaler_examples/finetune_hf_model/src/model_helper/customize/modeling_nnscaler_mixtral_4_42.py", line 358, in forward, value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) operator FwOp235-()(name=add, inputs=(t2029(p708,(1, 8192, 4096),d(),v(0/1)), t2063(p783,(1, 8192, 4096),d(),v(0/1))), outputs=(t2064(p785,(1, 8192, 4096),d(),v(0/1)),)) has 3 partitions, importance ratio 0.150 at File "/home/yizhu1/ts_dev/Fairseq/nnscaler_examples/finetune_hf_model/src/model_helper/customize/modeling_nnscaler_mixtral_4_42.py", line 831, in forward, hidden_states = residual + hidden_states operator FwOp160-()(name=add, inputs=(t1932(p457,(1, 8192, 4096),d(),v(0/1)), t1966(p532,(1, 8192, 4096),d(),v(0/1))), outputs=(t1967(p534,(1, 8192, 4096),d(),v(0/1)),)) has 3 partitions, importance ratio 0.150 at File "/home/yizhu1/ts_dev/Fairseq/nnscaler_examples/finetune_hf_model/src/model_helper/customize/modeling_nnscaler_mixtral_4_42.py", line 831, in forward, hidden_states = residual + hidden_states operator FwOp85-()(name=add, inputs=(t1835(p206,(1, 8192, 4096),d(),v(0/1)), t1869(p281,(1, 8192, 4096),d(),v(0/1))), outputs=(t1870(p283,(1, 8192, 4096),d(),v(0/1)),)) has 3 partitions, importance ratio 0.150 at File "/home/yizhu1/ts_dev/Fairseq/nnscaler_examples/finetune_hf_model/src/model_helper/customize/modeling_nnscaler_mixtral_4_42.py", line 831, in forward, hidden_states = residual + hidden_states ``` which means that constrain the partition space of embedding and residual add can reduce a large search space parity check & unit test passed
add nested output support
…le workers minitrainer: fix bug when running compile with multiple workers
parallel module: decouple from Program()
…oint If a tensor is a non-persistent buffer, we will check its existence when loading from merged checkpoint.
Add profiling capabilities to the CLI, allowing users to monitor CPU and CUDA activities during training.
Implement overlapped scheduler by: 1. add cuda stream context for each segment/adapter/reducer 2. add cuda event wait/record for each segment 3. refine sched code generation for stream/event config.
This PR introduces a “no-grad-reduce” annotation mechanism for custom op shape annotations so that, for specific partition identifiers, nnscaler can skip inserting gradient all-reduce adapters (avoiding incorrect or redundant reductions). This is done by extending ShapeAnno parsing to support : / modifiers (and '/' as a shortcut) to control gradient-reduction behavior during partitioning.
* Merge ring attention implementation into main branch * remove sink since it is no longer needed * remove cp_ranks * add tests
Introduce support for the grad_dtype attribute in parameters, enhancing flexibility in gradient precision management. --------- Co-authored-by: zyeric <cheerforwhy@gmail.com>
nns supports 2.11
#36) * [Fix] zero parameter-level sharding in Reducer and add corresponding unit tests * reduce the change, add more tests * refine message --------- Co-authored-by: Xu Weijiang <weijiangxu@microsoft.com>
…_data (#40) This pull request updates the broadcast_mixed_data function in nnscaler/utils.py to clarify and correct how the source rank and current rank are determined during distributed data broadcasting. The main focus is on ensuring the source rank is always interpreted in the context of the global process group, regardless of any custom group passed to the function. Clarification and correction of rank handling: The docstring for the src_rank argument is updated to specify that it refers to the global rank, not the rank within any custom process group, and that the source rank is always based on the global process group. The function now always retrieves the global rank using torch.distributed.get_rank() without a group argument, ensuring consistent behavior regardless of the group parameter.
…38) This pull request improves the handling of replicated tensors and gradient reduction logic in fn policy. It introduces a more precise distinction between replicated tensors that require gradient all-reduce and those that do not, so we can add multiref more precisely.
Correct the ordering of dataloaders in the pipeline to ensure consistent behavior across different ranks. This change addresses a bug that caused unexpected behavior due to incorrect insertion order.
* feat: annotate partition_descs with fqn and op in saved plan
Add fqn (fully qualified module name) and op (operator signature) fields
to each partition_descs entry in the saved plan JSON for improved
readability. Each entry is now a compact one-line dict:
{"cid": 11, "partition": [[[0, 0], 2]], "fqn": "...", "op": "..."}
The plan loader is backward compatible with old-format plans that use
[cid, desc] list entries.
* refactor: simplify _write_plan_json with placeholder approach
* perf: use single-pass regex replacement in _write_plan_json
* refactor: make cid2node required in to_json methods
- Remove cid2node=None fallback in TensorParallelDesc.to_json
- Add cid2node parameter to SPMDSearchOutput.to_json
- Make cid2node required in PipelineParallelDesc and PipelineSearchOutput
- Add type hints for cid2node: Dict[int, IRFwOperation]
* refactor: make cid2node optional, guard empty markers in _write_plan_json
- Make cid2node Optional[Dict[int, IRFwOperation]] with default None
- Skip fqn/op fields when cid2node is None
- Use direct import instead of TYPE_CHECKING
- Guard regex substitution in _write_plan_json for empty markers
* test(autodist): adapt test_follow helpers to dict-shaped partition_descs
partition_descs entries are now annotated dicts with cid/partition/fqn/op
(see commit d4f8fc2). Project them back to (cid, partition) tuples so the
existing expected_out tables in test_follow_attention and
test_solver_data_parallel continue to apply.
…#46) Fixes an incorrect bucket “chunkability” assertion in the reducer when using ZeRO-1 with zero_ngroups > 1, by validating divisibility against the ZeRO subgroup size (the actual sharding domain) rather than the full reducer group size. --------- Co-authored-by: Xu Weijiang <weijiangxu@microsoft.com>
… in pipeline. (#48) Adds a new configuration switch for fn policy to control whether pipeline-parallel stages insert multiref for shared parameters in pipeline
This PR adds pipeline-parallel support for returning and transferring non-tensor IRObjects by introducing runtime collectives (move_object, broadcast_object) and corresponding IR adapter primitives, then updating graph staging/adapter generation and tests to validate the new behavior.
This PR adds first-class support for segment-level pre/post hooks so callers can inject custom behavior immediately before/after each segment execution (forward and backward) in the generated pipeline schedule code.
Ensure the device_id is specified in the init_process_group to prevent warnings during barrier synchronization in distributed training.
Enhances ZeRO(1) parameter-level sharding behavior when a bucket has fewer parameters than the ZeRO subgroup size by allowing padding (so some ranks get empty shards), and updates related utilities/tests to validate the new behavior.
#42) * [Feat] Add return_lse parameter to attention functions for enhanced output options * [Fix] Update return_lse parameter to accept boolean values in attention functions * [Fix] Update output annotation in flash_attention_anno to reflect return_lse parameter changes * Avoid positional return_lse lookup in attention annotations * remove "raise RuntimeError for call_flash_attn_cute_varlen_func with return_lse = True" * [Feat] Enhance LSE handling in attention functions and add tests for return_lse behavior
This PR introduces an opt-in path to generate weight reducers for replicated weights, carrying an nreplicas value through IR → codegen → runtime so gradients can be divided after reduction. It also expands the test suite to validate the new flag’s effects on generated code and training. The main goal of this feature is to make sure gradients cross ranks are exactly the same after each step.
Introduce functionality to save and load non-persistent buffer content, enhancing model checkpointing and resuming capabilities. This update allows for efficient handling of non-persistent buffers by utilizing a dedicated file for their data. Additionally, implement tests to ensure the correct behavior of the new features.
1. Add AsyncLogger 2. Add new log config format to support async (and add backward compatibility support for old format) --------- Co-authored-by: XU Weijiang <90586345+0xWJ@users.noreply.github.com> Co-authored-by: Xu Weijiang <weijiangxu@microsoft.com>
Refine the reducer generation logic and enhance test coverage.
Refining the logic for inputs and outputs and dispatching to support segment TP. This is to reduce the generated communication in PP In comparison, in previous version, the inputs/outputs of segments are always full tensors even when the segment only uses a portion of the tensors. That design can introduce unnecessary communication and further have negative impact on performance.
1. Fix tensor_splits tracking for new tensors in fn during inserting new operators (identity and multiref) 2. Optimize dataloader usage in pipeline: prefer to read directly from dataloader instead of transmitting from stage 0.
* [Refine] More Flexible flatten * refine code * refine code * refine test * refine code * fix ci * add more checks
adds a slow_fs resume option so that on slow/remote filesystems only global rank 0 reads the merged checkpoint, then broadcasts it to each node's leader (local rank 0), which then follows the existing in-node trimmed broadcast. Default is False, so existing behavior is preserved. Overall it's clean, well-documented, and correct.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
max_seqlen_qandmax_seqlen_karguments to ring, sliding-window, and zigzag varlen attention wrappers.cu_seqlens.return_lseinterface.Motivation
llm-trainnow passes configured maximum sequence lengths into attention wrappers, avoiding repeated runtime derivation and supporting stable FlashAttention compilation.Testing
python -m pytest -q tests/customized_ops/ring_attn/test_return_lse.py8 passed