mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-08-17 13:35:44 +08:00
[v1] add FSDPTurbo EP/EFSDP plugin for MoE training (#10676)
This commit is contained in:
Hazeldxq
GitHub
1
.github/workflows/docs.yml
vendored
1
.github/workflows/docs.yml
vendored
@@ -66,6 +66,7 @@ jobs:
|
|||||||
path: docs/_build/html
|
path: docs/_build/html
|
||||||
|
|
||||||
deploy:
|
deploy:
|
||||||
|
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
||||||
environment:
|
environment:
|
||||||
name: github-pages
|
name: github-pages
|
||||||
url: ${{ steps.deployment.outputs.page_url }}
|
url: ${{ steps.deployment.outputs.page_url }}
|
||||||
|
|||||||
3
.github/workflows/tests_npu.yml
vendored
3
.github/workflows/tests_npu.yml
vendored
@@ -8,6 +8,7 @@ on:
|
|||||||
paths:
|
paths:
|
||||||
- "**/*.py"
|
- "**/*.py"
|
||||||
- "pyproject.toml"
|
- "pyproject.toml"
|
||||||
|
- "requirements/fsdpturbo.txt"
|
||||||
- "Makefile"
|
- "Makefile"
|
||||||
- ".github/workflows/*.yml"
|
- ".github/workflows/*.yml"
|
||||||
pull_request:
|
pull_request:
|
||||||
@@ -16,6 +17,7 @@ on:
|
|||||||
paths:
|
paths:
|
||||||
- "**/*.py"
|
- "**/*.py"
|
||||||
- "pyproject.toml"
|
- "pyproject.toml"
|
||||||
|
- "requirements/fsdpturbo.txt"
|
||||||
- "Makefile"
|
- "Makefile"
|
||||||
- ".github/workflows/*.yml"
|
- ".github/workflows/*.yml"
|
||||||
|
|
||||||
@@ -68,6 +70,7 @@ jobs:
|
|||||||
uv pip install -e .
|
uv pip install -e .
|
||||||
uv pip install -r requirements/npu.txt
|
uv pip install -r requirements/npu.txt
|
||||||
uv pip install -r requirements/dev.txt
|
uv pip install -r requirements/dev.txt
|
||||||
|
uv pip install --no-deps -r requirements/fsdpturbo.txt
|
||||||
|
|
||||||
- name: Install node
|
- name: Install node
|
||||||
run: |
|
run: |
|
||||||
|
|||||||
173
docs/en/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal file
173
docs/en/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal file
@@ -0,0 +1,173 @@
|
|||||||
|
# FSDPTurbo EP/EFSDP and LlamaFactory FSDP2/CP Design
|
||||||
|
|
||||||
|
Chinese version: [FSDPTurbo EP/EFSDP 与 LlamaFactory FSDP2/CP 设计说明](../../../zh/advanced/distributed/fsdpturbo-ep-efsdp.md)
|
||||||
|
|
||||||
|
This document describes the current implementation of the `fsdpturbo` distributed plugin. Its core principle is a clear separation of responsibilities:
|
||||||
|
|
||||||
|
- FSDPTurbo owns expert parallelism (EP), expert parameter sharding (EFSDP), and device operator registration.
|
||||||
|
- LlamaFactory owns process initialization, the base DeviceMesh, outer FSDP2, CP, model initialization, and weight loading.
|
||||||
|
- The LlamaFactory integration layer combines the two parameter layouts and handles gradient norms across meshes.
|
||||||
|
|
||||||
|
## 1. Configuration Boundaries
|
||||||
|
|
||||||
|
Common parallel topology belongs to `TrainingArguments`, while FSDPTurbo-only settings remain in `dist_config`:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
cp_size: 1
|
||||||
|
|
||||||
|
dist_config:
|
||||||
|
name: fsdpturbo
|
||||||
|
ep_size: 16
|
||||||
|
ep_dispatcher: eager
|
||||||
|
```
|
||||||
|
|
||||||
|
The fields used by the minimal example have the following responsibilities:
|
||||||
|
|
||||||
|
- `ep_size`: expert-parallel group size.
|
||||||
|
- `ep_dispatcher`: FSDPTurbo EP dispatcher, which defaults to `eager`.
|
||||||
|
|
||||||
|
`dp_size`, `cp_size`, `cp_mode`, `mp_replicate_size`, `mp_shard_size`, and `dist_timeout` are common topology fields and therefore remain at the top level. `dist_config` is parsed strictly as `FSDPTurboParams`; putting a common topology field inside it is rejected instead of being silently ignored.
|
||||||
|
|
||||||
|
The top-level training option `bf16` controls FSDPTurbo parameter storage and compute dtype. The backend casts the model before FSDP materialization, so `ModelEngine` does not need to read distributed-backend configuration.
|
||||||
|
|
||||||
|
The following advanced fields are optional and are therefore omitted from the minimal YAML example above:
|
||||||
|
|
||||||
|
- `fsdp_ignored_modules`: additional modules excluded from the outer LlamaFactory FSDP2 path. Expert parameters selected by the model spec are automatically added to the ignored set by the integration layer, so normal configurations do not need to repeat them here.
|
||||||
|
- `hook_modules`: optional module patterns for FSDPTurbo EFSDP hooks. The default is an empty list.
|
||||||
|
- `fsdp_implementation`: the FSDPTurbo EFSDP implementation, either `native` or `custom`. The default is `native`.
|
||||||
|
|
||||||
|
The model spec determines the EFSDP targets. Non-expert parameters such as attention, embeddings, and the LM head do not enter the FSDPTurbo EFSDP plan. They remain managed by the outer LlamaFactory FSDP2 layer.
|
||||||
|
|
||||||
|
Model-specific module paths and preparation logic are managed exclusively by the `FSDPTurboEPModelSpec` registry. Built-in specs currently cover `qwen3_moe` and `qwen3_5_moe`; unregistered models fail with an explicit error. `ep_modules` and `ep_fsdp_modules` are not YAML options, and strict parameter parsing rejects them to prevent configuration from drifting away from the actual model structure.
|
||||||
|
|
||||||
|
## 2. Mesh Initialization
|
||||||
|
|
||||||
|
LlamaFactory's `DistributedInterface` initializes only its existing model and data meshes. It is unaware of EP and EFSDP and does not expose an extra mesh registration interface for distributed plugins. The FSDPTurbo expert topology is independently created and owned by `FSDPTurboParallelState` in the plugin module:
|
||||||
|
|
||||||
|
```text
|
||||||
|
run_sft / run_dpo / run_rm
|
||||||
|
-> DistributedInterface(training_args)
|
||||||
|
-> initialize LlamaFactory model/data meshes
|
||||||
|
-> DistributedPlugin("fsdpturbo").shard_model(...)
|
||||||
|
-> FSDPTurboFSDP2Engine.__init__()
|
||||||
|
-> FSDPTurboParallelState.initialize()
|
||||||
|
-> initialize and retain the expert parent mesh and submeshes
|
||||||
|
```
|
||||||
|
|
||||||
|
`FSDPTurboParallelState` creates a four-dimensional expert parent mesh:
|
||||||
|
|
||||||
|
```text
|
||||||
|
(edp, efsdp, ep, expert_cp)
|
||||||
|
```
|
||||||
|
|
||||||
|
Its current dimensions are calculated as follows:
|
||||||
|
|
||||||
|
```text
|
||||||
|
dp_size = world_size / cp_size
|
||||||
|
ep_fsdp_size = dp_size / ep_size
|
||||||
|
edp_size = dp_size / (ep_size * ep_fsdp_size)
|
||||||
|
mesh_shape = (edp_size, ep_fsdp_size, ep_size, cp_size)
|
||||||
|
```
|
||||||
|
|
||||||
|
The state object retains `edp_mesh`, `efsdp_mesh`, `ep_mesh`, and `expert_cp_mesh`. Model sharding and gradient norm logic inside the plugin read expert communication domains from this state, while other LlamaFactory backends do not need to implement or know about these interfaces. Initialization validates that `ep_size` is positive and divides `dp_size`; repeated initialization also rejects topology changes.
|
||||||
|
|
||||||
|
## 3. Model Sharding Order
|
||||||
|
|
||||||
|
The wrapping order must remain "expert side first, outer FSDP2 second":
|
||||||
|
|
||||||
|
```text
|
||||||
|
DistributedPlugin("fsdpturbo")
|
||||||
|
-> FSDPTurboFSDP2Engine.shard_model(model)
|
||||||
|
-> prepare_model_ep(model)
|
||||||
|
-> expert_parallelize_modules(model, ep_mesh, ep_plan)
|
||||||
|
-> expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)
|
||||||
|
-> collect expert parameters as ignored_params
|
||||||
|
-> FSDP2Engine.prepare_model(model, ignored_params=...)
|
||||||
|
-> apply outer fully_shard to the remaining Transformer Layers and root module
|
||||||
|
```
|
||||||
|
|
||||||
|
This prevents the same expert parameter from being managed by both EFSDP and outer FSDP2. Outer FSDP2 continues to reuse LlamaFactory's model initialization, checkpoint, and save flows.
|
||||||
|
|
||||||
|
The LlamaFactory integration layer accepts `eager`, `fused`, `mc2`, and `domino` and forwards the selected value unchanged to FSDPTurbo. Their implementation boundaries and current validation status differ:
|
||||||
|
|
||||||
|
| Dispatcher | Main path | Additional requirements | Validation in this PR |
|
||||||
|
| --- | --- | --- | --- |
|
||||||
|
| `eager` | Uses PyTorch implementations of permute, unpermute, and grouped matmul while tensors remain on the current accelerator, with standard AllToAll for token dispatch and combine | Minimal dependencies; serves as the reference implementation | End-to-end numerical and performance validation completed on Ascend A3 |
|
||||||
|
| `fused` | Keeps the same AllToAll topology while replacing permute, unpermute, and grouped matmul with device-fused operators | Requires matching device operators, dtypes, and layouts; local operators may fall back to eager when an expert receives no tokens | End-to-end numerical and performance validation completed on Ascend A3 |
|
||||||
|
| `mc2` | Uses dedicated operators that fuse AllToAllV with grouped matmul to reduce intermediate communication-computation overhead | Requires the MC2 NPU operators, an HCCL communicator, and their shape and dtype constraints | Implemented by FSDPTurbo but not validated end to end in this PR |
|
||||||
|
| `domino` | Splits the first dimension of the expert-module input into two slices and uses a separate communication stream and events to overlap AllToAll with expert computation | Requires asynchronous stream/event support and enough token work in both slices to amortize scheduling overhead | Implemented by FSDPTurbo but not validated end to end in this PR |
|
||||||
|
|
||||||
|
Only `eager` and `fused` are validated here because they cover the reference path and the commonly used A3 device-fused path, respectively, and therefore isolate and establish the correctness of the EP/EFSDP integration between LlamaFactory and FSDPTurbo. The current experiment matrix was not extended to `mc2` and `domino`: they add operator, communication-scheduling, and input-shape constraints that require separate numerical comparisons, long-run stability tests, and profiler analysis. They are accepted configuration choices, but the results in this PR should not be interpreted as evidence that they have reached the same stability, numerical, or performance level.
|
||||||
|
|
||||||
|
## 4. FSDPTurbo Dependency Entry Points
|
||||||
|
|
||||||
|
LlamaFactory imports each required object directly from the module that defines it:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import (
|
||||||
|
expert_fully_shard_modules,
|
||||||
|
)
|
||||||
|
from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules
|
||||||
|
from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig
|
||||||
|
from fsdp_turbo.utils.str_match import module_name_match
|
||||||
|
```
|
||||||
|
|
||||||
|
The imports occur inside `prepare_model_ep()`, so other distributed backends remain importable when FSDPTurbo is not installed. They intentionally bypass aggregate exports from `fsdp_turbo.distributed.__init__` to avoid extra dependencies and potential import cycles during package initialization.
|
||||||
|
|
||||||
|
## 5. Gradient Norms
|
||||||
|
|
||||||
|
Outer and expert parameters can belong to different DTensor meshes and therefore cannot be passed together to a single standard `clip_grad_norm_()` call. The `fsdpturbo` plugin groups parameters by their owning mesh and computes local p-power sums:
|
||||||
|
|
||||||
|
- Non-expert parameters are reduced over the DP and CP groups.
|
||||||
|
- Expert parameters are reduced over the EFSDP, EP, and expert-CP groups retained by `FSDPTurboParallelState`.
|
||||||
|
- After the global norm is assembled, the same clipping coefficient is applied to every local gradient.
|
||||||
|
|
||||||
|
A zero-gradient warmup runs during startup so that the required collectives are initialized before training begins. This is currently a backend-specific implementation for `fsdpturbo`; other backends retain their existing gradient norm paths until the upstream distributed plugin interface is decoupled.
|
||||||
|
|
||||||
|
## 6. Weight Loading
|
||||||
|
|
||||||
|
LlamaFactory retains the `init_on_meta` and safetensors loading flow. The parent `FSDP2Engine` loader dynamically invokes the FSDPTurbo engine override through `self._copy_weights(...)`, so the method is not dead code. It supports DTensors with multiple `Shard` placements by calculating the local slice for the current rank along each mesh dimension in sequence. Model save and checkpoint interfaces continue to reuse the LlamaFactory FSDP2 implementation.
|
||||||
|
|
||||||
|
## 7. Kernel Plugin
|
||||||
|
|
||||||
|
FLA operators do not belong in the distributed configuration. Operator selection is handled through an independent `kernel_config`:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
kernel_config:
|
||||||
|
name: auto, flash-linear-attention
|
||||||
|
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
|
||||||
|
chunk_size: 32
|
||||||
|
```
|
||||||
|
|
||||||
|
The call path is:
|
||||||
|
|
||||||
|
```text
|
||||||
|
ModelEngine
|
||||||
|
-> apply_kernels("auto, flash-linear-attention")
|
||||||
|
-> accelerator-specific LlamaFactory auto kernels
|
||||||
|
-> KernelPlugin("flash-linear-attention").apply(...)
|
||||||
|
-> fsdp_turbo.ops.get_op()
|
||||||
|
-> FSDPTurbo device operator registry
|
||||||
|
-> fsdp_turbo.utils.patch.patch_model_members()
|
||||||
|
-> FLA backend implementation
|
||||||
|
```
|
||||||
|
|
||||||
|
`chunk_size` accepts `16`, `32`, and `64`, with a default of `64`. The kernel plugin and distributed plugin are independent. `name: flash-linear-attention` installs only the selected FLA operators. The comma-separated `name: auto, flash-linear-attention` form composes LlamaFactory's accelerator-specific automatic kernels with the FLA plugin before distributed sharding. LlamaFactory owns the operator-to-model-attribute mapping and `chunk_size` binding; FSDPTurbo owns device operator registration, selection, and generic callable patching. FLA stays explicit because it has optional external dependencies and is not part of the built-in `auto` set. FSDPTurbo subsequently replaces the target expert module's `forward`, so the final expert execution path is selected by `ep_dispatcher`; an MoE kernel applied during the auto stage is not retained as a separate second expert execution path.
|
||||||
|
|
||||||
|
## 8. CP Runtime Constraints and Validation Scope
|
||||||
|
|
||||||
|
When `init_on_meta` constructs the model, it must propagate `attn_implementation` in the same way as the `from_pretrained` path. Otherwise, the model falls back to a non-FlashAttention implementation and Ulysses CP cannot start. Before calling Hugging Face FlashAttention, Ulysses reconstructs the global attention mask. Only two-dimensional position IDs participate in packed-sequence detection. Multi-axis position IDs such as Qwen3.5 mRoPE have already been consumed by rotary embedding and must not be passed to the FlashAttention packed-sequence detection logic.
|
||||||
|
|
||||||
|
The current implementation has completed the following BF16 AdamW full SFT validations with Qwen3.5-35B-A3B on Atlas 900 A3 SuperPoD and Atlas 950 SuperPoD systems. This revalidation used FSDPTurbo `0e96fbc`. The A3 environment used CANN 9.0.0, PyTorch 2.7.1, and torch-npu 2.7.1.post4; the A5 environment used CANN 9.1.0-beta.3, PyTorch 2.10.0, and torch-npu 2.10.0.post2. Performance is calculated from the step 1 and step 100 log timestamps and excludes initialization and compilation before the first step as well as model saving after training:
|
||||||
|
|
||||||
|
| Machine | CP | EP | EFSDP | Checkpoint | Kernel / Dispatcher | Steps | Loss (first -> last) | Performance | Result |
|
||||||
|
| --- | ---: | ---: | ---: | --- | --- | ---: | --- | ---: | --- |
|
||||||
|
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | Off | FLA (chunk size 16) / eager | 100 | 1.3361 -> 0.0793 | 2.51 s/it | Passed and saved |
|
||||||
|
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | Off | FLA (chunk size 16) / fused | 100 | 1.3354 -> 0.1179 | 2.17 s/it | Passed and saved |
|
||||||
|
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | Off | auto + FLA (chunk size 64) / fused | 100 | 1.8114 -> 0.5260 | 7.65 s/it | Passed and saved |
|
||||||
|
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | Off | auto + FLA (chunk size 64) / eager | 100 | 1.8095 -> 0.5596 | 5.88 s/it | Passed and saved |
|
||||||
|
| Atlas 950 SuperPoD | 1 | 8 | 1 | Off | no kernel plugin configured / eager | 100 | 1.3575 -> 0.4439 | 2.68 s/it | Passed and saved |
|
||||||
|
|
||||||
|
Loss and gradient norm remained finite in all five runs, and every run completed 100 steps and model saving. With the same partition, the per-step loss correlation between eager and fused was 0.997 for EP16 and 0.977 for CP2/EP4/EFSDP2, which indicates consistent optimization trajectories. The performance effect depends on the partition: fused was about 13% faster than eager with EP16, but about 30% slower after adding CP and EFSDP. Fused therefore should not be treated as the default optimum for every mesh.
|
||||||
|
|
||||||
|
The EP16 runs used global batch 16 and cutoff length 256. The CP2 runs used global batch 8 and cutoff length 128. The A5 run used global batch 8 and cutoff length 256. The first-to-last loss validates convergence within each run; absolute loss values across different partition groups should not be used directly as a numerical-equivalence conclusion.
|
||||||
@@ -38,6 +38,7 @@ LlamaFactory Docs
|
|||||||
advanced/distributed/fsdp
|
advanced/distributed/fsdp
|
||||||
advanced/distributed/deepspeed
|
advanced/distributed/deepspeed
|
||||||
advanced/distributed/parallel-dp-tp-ep-sp-cp
|
advanced/distributed/parallel-dp-tp-ep-sp-cp
|
||||||
|
advanced/distributed/fsdpturbo-ep-efsdp
|
||||||
advanced/custom-kernels/triton
|
advanced/custom-kernels/triton
|
||||||
advanced/custom-kernels/fused-operators
|
advanced/custom-kernels/fused-operators
|
||||||
|
|
||||||
|
|||||||
215
docs/zh/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal file
215
docs/zh/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal file
@@ -0,0 +1,215 @@
|
|||||||
|
# FSDPTurbo EP/EFSDP 与 LlamaFactory FSDP2/CP 设计说明
|
||||||
|
|
||||||
|
English version: [FSDPTurbo EP/EFSDP and LlamaFactory FSDP2/CP Design](../../../en/advanced/distributed/fsdpturbo-ep-efsdp.md)
|
||||||
|
|
||||||
|
本文描述 `fsdpturbo` distributed plugin 的当前实现。核心原则是保持两侧职责清晰：
|
||||||
|
|
||||||
|
- FSDPTurbo 负责专家并行（EP）、专家参数分片（EFSDP）和设备算子注册。
|
||||||
|
- LlamaFactory 负责进程初始化、基础 DeviceMesh、外层 FSDP2、CP、模型初始化与权重加载。
|
||||||
|
- LlamaFactory 的集成层负责把两套参数布局组合起来，并处理跨 Mesh 的梯度范数。
|
||||||
|
|
||||||
|
## 1. 配置边界
|
||||||
|
|
||||||
|
公共并行拓扑放在 `TrainingArguments` 顶层，FSDPTurbo 私有参数保留在 `dist_config`：
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
cp_size: 1
|
||||||
|
|
||||||
|
dist_config:
|
||||||
|
name: fsdpturbo
|
||||||
|
ep_size: 16
|
||||||
|
ep_dispatcher: eager
|
||||||
|
```
|
||||||
|
|
||||||
|
最小示例中的字段职责如下：
|
||||||
|
|
||||||
|
- `ep_size`：专家并行组大小。
|
||||||
|
- `ep_dispatcher`：FSDPTurbo EP dispatcher，默认为 `eager`。
|
||||||
|
|
||||||
|
`dp_size`、`cp_size`、`cp_mode`、`mp_replicate_size`、`mp_shard_size` 和 `dist_timeout`
|
||||||
|
属于公共拓扑字段，继续放在顶层。`dist_config` 会被严格解析为 `FSDPTurboParams`；如果把公共拓扑
|
||||||
|
字段误放进去，会直接报错，而不是静默忽略。
|
||||||
|
|
||||||
|
顶层训练参数 `bf16` 同时控制 FSDPTurbo 的参数存储和计算 dtype。backend 会在 FSDP materialization
|
||||||
|
前完成模型 dtype 转换，因此 `ModelEngine` 不需要读取 distributed backend 配置。
|
||||||
|
|
||||||
|
以下高级字段为可选项，因此没有写入上面的最小 YAML 示例：
|
||||||
|
|
||||||
|
- `fsdp_ignored_modules`：额外排除在 LlamaFactory 外层 FSDP2 之外的模块。模型规格选中的专家参数
|
||||||
|
会被集成层自动加入忽略集合，普通配置无需重复填写。
|
||||||
|
- `hook_modules`：FSDPTurbo EFSDP hook 的可选模块模式，默认为空列表。
|
||||||
|
- `fsdp_implementation`：FSDPTurbo EFSDP 实现，可选 `native` 或 `custom`，默认为 `native`。
|
||||||
|
|
||||||
|
EFSDP 的目标由模型规格决定。Attention、Embedding、LM Head 等非专家参数不进入 FSDPTurbo
|
||||||
|
EFSDP plan，而是继续由 LlamaFactory 外层 FSDP2 管理。
|
||||||
|
|
||||||
|
模型相关的模块路径和准备逻辑统一由 `FSDPTurboEPModelSpec` 注册表管理。当前内置 `qwen3_moe`
|
||||||
|
和 `qwen3_5_moe`；未注册的模型会明确报错。`ep_modules` 和 `ep_fsdp_modules` 不属于 YAML
|
||||||
|
接口，严格参数解析会拒绝这两个字段，避免用户配置与模型实际结构失配。
|
||||||
|
|
||||||
|
## 2. Mesh 初始化
|
||||||
|
|
||||||
|
LlamaFactory 的 `DistributedInterface` 只初始化自身原有的 model/data mesh。它不感知 EP、EFSDP，
|
||||||
|
也不为 distributed plugin 提供额外 mesh 注册接口。FSDPTurbo 的专家拓扑由插件文件内的
|
||||||
|
`FSDPTurboParallelState` 独立创建和持有：
|
||||||
|
|
||||||
|
```text
|
||||||
|
run_sft / run_dpo / run_rm
|
||||||
|
-> DistributedInterface(training_args)
|
||||||
|
-> 初始化 LlamaFactory model/data mesh
|
||||||
|
-> DistributedPlugin("fsdpturbo").shard_model(...)
|
||||||
|
-> FSDPTurboFSDP2Engine.__init__()
|
||||||
|
-> FSDPTurboParallelState.initialize()
|
||||||
|
-> 初始化并保存 expert parent mesh 及其子 mesh
|
||||||
|
```
|
||||||
|
|
||||||
|
`FSDPTurboParallelState` 创建专家侧四维父 Mesh：
|
||||||
|
|
||||||
|
```text
|
||||||
|
(edp, efsdp, ep, expert_cp)
|
||||||
|
```
|
||||||
|
|
||||||
|
当前尺寸计算为：
|
||||||
|
|
||||||
|
```text
|
||||||
|
dp_size = world_size / cp_size
|
||||||
|
ep_fsdp_size = dp_size / ep_size
|
||||||
|
edp_size = dp_size / (ep_size * ep_fsdp_size)
|
||||||
|
mesh_shape = (edp_size, ep_fsdp_size, ep_size, cp_size)
|
||||||
|
```
|
||||||
|
|
||||||
|
状态对象保存 `edp_mesh`、`efsdp_mesh`、`ep_mesh` 和 `expert_cp_mesh`。插件内部模型切分和梯度范数
|
||||||
|
都从这个状态对象读取专家通信域；LlamaFactory 其他 backend 不需要实现或感知这些接口。状态初始化
|
||||||
|
会校验 `ep_size` 为正数且能够整除 `dp_size`，重复初始化时也会拒绝拓扑发生变化。
|
||||||
|
|
||||||
|
## 3. 模型切分顺序
|
||||||
|
|
||||||
|
模型包装顺序必须保持为“专家侧优先，外层 FSDP2 随后”：
|
||||||
|
|
||||||
|
```text
|
||||||
|
DistributedPlugin("fsdpturbo")
|
||||||
|
-> FSDPTurboFSDP2Engine.shard_model(model)
|
||||||
|
-> prepare_model_ep(model)
|
||||||
|
-> expert_parallelize_modules(model, ep_mesh, ep_plan)
|
||||||
|
-> expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)
|
||||||
|
-> 收集专家参数作为 ignored_params
|
||||||
|
-> FSDP2Engine.prepare_model(model, ignored_params=...)
|
||||||
|
-> 对剩余 Transformer Layer 和根模块执行 outer fully_shard
|
||||||
|
```
|
||||||
|
|
||||||
|
这样可以避免同一专家参数同时被 EFSDP 和外层 FSDP2 管理。外层 FSDP2 仍复用 LlamaFactory
|
||||||
|
原有的初始化、checkpoint 和保存流程。
|
||||||
|
|
||||||
|
LlamaFactory 集成层接受 `eager`、`fused`、`mc2` 和 `domino`，并将选项原样传给
|
||||||
|
FSDPTurbo。这四种模式的实现边界和当前验证状态不同：
|
||||||
|
|
||||||
|
| Dispatcher | 主要路径 | 额外要求 | 本 PR 验证状态 |
|
||||||
|
| --- | --- | --- | --- |
|
||||||
|
| `eager` | 使用 PyTorch 实现 permute、unpermute 和 grouped matmul，张量仍在当前加速设备上，通过标准 AllToAll 完成 token dispatch/combine | 依赖最少，用作参考实现 | 已在 A3 上完成精度和性能验证 |
|
||||||
|
| `fused` | 保持相同的 AllToAll 拓扑，将 permute、unpermute 和 grouped matmul 切换为设备融合算子 | 需要对应的设备算子、dtype 和 layout 支持；存在空专家时可回退到 eager 局部算子 | 已在 A3 上完成精度和性能验证 |
|
||||||
|
| `mc2` | 使用专用算子融合 AllToAllV 和 grouped matmul，减少通信与计算之间的中间开销 | 依赖 MC2 NPU 算子、HCCL communicator 以及对应的 shape/dtype 约束 | FSDPTurbo 提供实现，本 PR 未做端到端验证 |
|
||||||
|
| `domino` | 将专家模块输入的第一维分成两片，使用独立通信流和 event 重叠 AllToAll 与专家计算 | 需要异步 stream/event 支持，且两个分片都要有足够的 token 工作量才能覆盖调度开销 | FSDPTurbo 提供实现，本 PR 未做端到端验证 |
|
||||||
|
|
||||||
|
当前只验证 `eager` 和 `fused`，是因为它们分别覆盖参考实现和 A3 常用设备融合路径，可用于隔离并验证
|
||||||
|
LlamaFactory 与 FSDPTurbo 之间的 EP/EFSDP 集成正确性。本次实验矩阵没有继续扩展到 `mc2` 和
|
||||||
|
`domino`：它们还引入了额外的算子、通信调度和输入形状约束，需要独立比较数值、长步稳定性和 profiler 结果。
|
||||||
|
因此，它们在配置接口上可选，但不应从本 PR 的实验结果推断为已达到相同的稳定性、精度或性能水平。
|
||||||
|
|
||||||
|
## 4. FSDPTurbo 依赖入口
|
||||||
|
|
||||||
|
LlamaFactory 从各功能的定义模块直接导入所需对象：
|
||||||
|
|
||||||
|
```python
|
||||||
|
from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import (
|
||||||
|
expert_fully_shard_modules,
|
||||||
|
)
|
||||||
|
from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules
|
||||||
|
from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig
|
||||||
|
from fsdp_turbo.utils.str_match import module_name_match
|
||||||
|
```
|
||||||
|
|
||||||
|
导入发生在 `prepare_model_ep()` 内，因此没有安装 FSDPTurbo 时，其他 distributed backend 仍可正常导入。
|
||||||
|
这里不通过 `fsdp_turbo.distributed.__init__` 聚合导出，避免 package 初始化期间的额外依赖和潜在循环导入。
|
||||||
|
|
||||||
|
## 5. 梯度范数
|
||||||
|
|
||||||
|
外层参数和专家参数可能属于不同 DTensor Mesh，不能直接放入一次标准 `clip_grad_norm_()`。
|
||||||
|
`fsdpturbo` plugin 按参数所属 Mesh 分组计算局部 p 次方和：
|
||||||
|
|
||||||
|
- 非专家参数沿 DP 和 CP group 汇总。
|
||||||
|
- 专家参数沿 `FSDPTurboParallelState` 保存的 EFSDP、EP 和 expert-CP group 汇总。
|
||||||
|
- 汇总得到全局范数后，对所有本地梯度应用同一个 clipping coefficient。
|
||||||
|
|
||||||
|
启动阶段会执行一次零梯度 warmup，使相关 collective 在正式训练前完成初始化。
|
||||||
|
当前这是 `fsdpturbo` backend 的专用实现；其他 backend 继续保留原有梯度范数路径，等待上游
|
||||||
|
distributed plugin 解耦后再统一公共接口。
|
||||||
|
|
||||||
|
## 6. 权重加载
|
||||||
|
|
||||||
|
LlamaFactory 保留 `init_on_meta` 和 safetensors 加载流程。父类 `FSDP2Engine` 的加载器通过
|
||||||
|
`self._copy_weights(...)` 动态调用 FSDPTurbo engine 的覆写实现，因此该方法不是未使用代码。
|
||||||
|
它支持包含多个 `Shard` placement 的 DTensor，按各 Mesh 维度依次计算当前 rank 对应的本地切片。
|
||||||
|
模型保存和 checkpoint 接口继续复用 LlamaFactory FSDP2 实现。
|
||||||
|
|
||||||
|
## 7. Kernel plugin
|
||||||
|
|
||||||
|
FLA 算子不属于 distributed config。算子选择通过独立的 `kernel_config` 完成：
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
kernel_config:
|
||||||
|
name: auto, flash-linear-attention
|
||||||
|
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
|
||||||
|
chunk_size: 32
|
||||||
|
```
|
||||||
|
|
||||||
|
调用链如下：
|
||||||
|
|
||||||
|
```text
|
||||||
|
ModelEngine
|
||||||
|
-> apply_kernels("auto, flash-linear-attention")
|
||||||
|
-> LlamaFactory 当前加速器对应的 auto kernels
|
||||||
|
-> KernelPlugin("flash-linear-attention").apply(...)
|
||||||
|
-> fsdp_turbo.ops.get_op()
|
||||||
|
-> FSDPTurbo device operator registry
|
||||||
|
-> fsdp_turbo.utils.patch.patch_model_members()
|
||||||
|
-> FLA backend implementation
|
||||||
|
```
|
||||||
|
|
||||||
|
`chunk_size` 当前支持 `16`、`32` 和 `64`，默认值为 `64`。Kernel plugin 与 distributed plugin
|
||||||
|
彼此独立。`name: flash-linear-attention` 只安装所选 FLA 算子；逗号分隔的
|
||||||
|
`name: auto, flash-linear-attention` 会在分布式切分前组合 LlamaFactory 当前加速器的 auto kernels
|
||||||
|
与 FLA plugin。LlamaFactory 负责算子名到模型属性的映射和 `chunk_size` 参数绑定；FSDPTurbo 负责设备
|
||||||
|
算子注册、选择和通用 callable patch。FLA 依赖可选的外部三方件，因此保持显式选择，不属于内置
|
||||||
|
`auto` 集合。FSDPTurbo
|
||||||
|
随后会替换目标专家模块的 `forward`，所以专家计算的最终路径由 `ep_dispatcher` 决定；auto 阶段
|
||||||
|
应用的 MoE kernel 不会作为独立的第二条专家执行路径保留下来。
|
||||||
|
|
||||||
|
## 8. CP 运行约束与验证范围
|
||||||
|
|
||||||
|
`init_on_meta` 构造模型时必须与 `from_pretrained` 路径一样传递 `attn_implementation`，否则模型会退回
|
||||||
|
非 FlashAttention 实现，Ulysses CP 无法启动。Ulysses 在调用 Hugging Face FlashAttention 前重建全局
|
||||||
|
attention mask；只有二维 position IDs 才参与 packed-sequence 检测。Qwen3.5 mRoPE 等多轴 position IDs
|
||||||
|
已经在 rotary embedding 中消费，不应传入 FlashAttention 的 packed-sequence 检测逻辑。
|
||||||
|
|
||||||
|
当前实现已在 Atlas 900 A3 SuperPoD 和 Atlas 950 SuperPoD 上用 Qwen3.5-35B-A3B 完成以下
|
||||||
|
BF16、AdamW full SFT 验证。本次重验证使用 FSDPTurbo `0e96fbc`；A3 环境为 CANN 9.0.0、
|
||||||
|
PyTorch 2.7.1 和 torch-npu 2.7.1.post4，A5 环境为 CANN 9.1.0-beta.3、PyTorch 2.10.0 和
|
||||||
|
torch-npu 2.10.0.post2。表中性能按第 1 步至第 100 步的日志时间戳计算，不包含首步前的初始化、
|
||||||
|
编译和训练后的模型保存时间：
|
||||||
|
|
||||||
|
| 机器型号 | CP | EP | EFSDP | Checkpoint | Kernel / Dispatcher | 步数 | Loss（首步 -> 末步） | 性能 | 结果 |
|
||||||
|
| --- | ---: | ---: | ---: | --- | --- | ---: | --- | ---: | --- |
|
||||||
|
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | 关闭 | FLA（chunk size 16）/ eager | 100 | 1.3361 -> 0.0793 | 2.51 s/it | 通过并完成保存 |
|
||||||
|
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | 关闭 | FLA（chunk size 16）/ fused | 100 | 1.3354 -> 0.1179 | 2.17 s/it | 通过并完成保存 |
|
||||||
|
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | 关闭 | auto + FLA（chunk size 64）/ fused | 100 | 1.8114 -> 0.5260 | 7.65 s/it | 通过并完成保存 |
|
||||||
|
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | 关闭 | auto + FLA（chunk size 64）/ eager | 100 | 1.8095 -> 0.5596 | 5.88 s/it | 通过并完成保存 |
|
||||||
|
| Atlas 950 SuperPoD | 1 | 8 | 1 | 关闭 | 未配置 kernel plugin / eager | 100 | 1.3575 -> 0.4439 | 2.68 s/it | 通过并完成保存 |
|
||||||
|
|
||||||
|
五组训练的 loss 和 grad norm 均保持有限，并完成 100 步及模型保存。同一切分下，EP16 eager/fused
|
||||||
|
的逐步 loss 相关系数为 0.997，CP2/EP4/EFSDP2 eager/fused 为 0.977，说明两种 dispatcher 的
|
||||||
|
优化轨迹一致。性能收益与切分有关：EP16 下 fused 比 eager 快约 13%，而加入 CP 和 EFSDP 后 fused
|
||||||
|
比 eager 慢约 30%，因此不能把 fused 视为所有 mesh 的默认最优选择。
|
||||||
|
|
||||||
|
EP16 两组使用 global batch 16 和 cutoff length 256；CP2 两组使用 global batch 8 和 cutoff length
|
||||||
|
128；A5 组使用 global batch 8 和 cutoff length 256。因此，首末 loss 用于验证各组自身的收敛趋势，
|
||||||
|
不同切分组之间的绝对 loss 不应直接作为精度等价结论。
|
||||||
@@ -38,6 +38,7 @@ LlamaFactory 文档
|
|||||||
advanced/distributed/fsdp
|
advanced/distributed/fsdp
|
||||||
advanced/distributed/deepspeed
|
advanced/distributed/deepspeed
|
||||||
advanced/distributed/parallel-dp-tp-ep-sp-cp
|
advanced/distributed/parallel-dp-tp-ep-sp-cp
|
||||||
|
advanced/distributed/fsdpturbo-ep-efsdp
|
||||||
advanced/custom-kernels/triton
|
advanced/custom-kernels/triton
|
||||||
advanced/custom-kernels/fused-operators
|
advanced/custom-kernels/fused-operators
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,32 @@
|
|||||||
|
model: Qwen/Qwen3.5-35B-A3B
|
||||||
|
model_class: llm
|
||||||
|
|
||||||
|
kernel_config:
|
||||||
|
name: auto, flash-linear-attention
|
||||||
|
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
|
||||||
|
chunk_size: 64
|
||||||
|
|
||||||
|
dist_config:
|
||||||
|
name: fsdpturbo
|
||||||
|
ep_size: 16
|
||||||
|
ep_dispatcher: eager
|
||||||
|
|
||||||
|
cp_size: 1
|
||||||
|
|
||||||
|
init_config:
|
||||||
|
name: init_on_meta
|
||||||
|
|
||||||
|
### data
|
||||||
|
train_dataset: data/v1_sft_demo.yaml
|
||||||
|
|
||||||
|
### training
|
||||||
|
output_dir: outputs/Qwen3.5-35B-A3B/full/sft
|
||||||
|
micro_batch_size: 1
|
||||||
|
cutoff_len: 256
|
||||||
|
learning_rate: 1.0e-4
|
||||||
|
bf16: true
|
||||||
|
max_steps: 5
|
||||||
|
|
||||||
|
### sample
|
||||||
|
sample_backend: hf
|
||||||
|
max_new_tokens: 128
|
||||||
2
requirements/fsdpturbo.txt
Normal file
2
requirements/fsdpturbo.txt
Normal file
@@ -0,0 +1,2 @@
|
|||||||
|
# Pin the FSDPTurbo API to a reproducible upstream main revision.
|
||||||
|
fsdp-turbo @ git+https://gitcode.com/Ascend/FSDPTurbo.git@d878dffdb1e0312dc098599f2b56810d6b592ee2
|
||||||
@@ -293,19 +293,24 @@ class BaseTrainer:
|
|||||||
if self._deepspeed_engine is not None:
|
if self._deepspeed_engine is not None:
|
||||||
# deepspeed: engine.step() already ran inside backward at the sync boundary
|
# deepspeed: engine.step() already ran inside backward at the sync boundary
|
||||||
grad_norm = self._deepspeed_engine.get_grad_norm()
|
grad_norm = self._deepspeed_engine.get_grad_norm()
|
||||||
|
else:
|
||||||
|
dist_name = self.args.dist_config.name if self.args.dist_config else None
|
||||||
|
if dist_name == "fsdpturbo":
|
||||||
|
from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
||||||
|
|
||||||
|
grad_norm = DistributedPlugin(dist_name).clip_grad_norm(self.model, self.args.max_grad_norm)
|
||||||
else:
|
else:
|
||||||
# FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a
|
# FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a
|
||||||
# per-rank local shard norm (global / sqrt(shard_size)): reported grad_norm then
|
# per-rank local shard norm. Materialize the true global norm before clipping.
|
||||||
# scales as 1/sqrt(dp_size) and the clip coefficient is applied per-shard. Reduce
|
|
||||||
# to the true global norm first, then clip with it.
|
|
||||||
grads = [p.grad for p in self.model.parameters() if p.grad is not None]
|
grads = [p.grad for p in self.model.parameters() if p.grad is not None]
|
||||||
total_norm = torch.nn.utils.get_total_norm(grads)
|
total_norm = torch.nn.utils.get_total_norm(grads)
|
||||||
if isinstance(total_norm, DTensor):
|
if isinstance(total_norm, DTensor):
|
||||||
# full_tensor all-reduces across the fsdp mesh (spans CP under default
|
# full_tensor all-reduces across the fsdp mesh (spans CP under default
|
||||||
# mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size).
|
# mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size).
|
||||||
total_norm = total_norm.full_tensor()
|
total_norm = total_norm.full_tensor()
|
||||||
# pass a Tensor: clip_grads_with_norm_ clamps max_norm / (total_norm + 1e-6).
|
torch.nn.utils.clip_grads_with_norm_(
|
||||||
torch.nn.utils.clip_grads_with_norm_(self.model.parameters(), self.args.max_grad_norm, total_norm)
|
self.model.parameters(), self.args.max_grad_norm, total_norm
|
||||||
|
)
|
||||||
grad_norm = total_norm.item()
|
grad_norm = total_norm.item()
|
||||||
|
|
||||||
if not torch.isfinite(torch.tensor(grad_norm)): # type: ignore # pyright: ignore [reportUnknownReturnType]
|
if not torch.isfinite(torch.tensor(grad_norm)): # type: ignore # pyright: ignore [reportUnknownReturnType]
|
||||||
@@ -363,7 +368,7 @@ class BaseTrainer:
|
|||||||
|
|
||||||
def save_model(self) -> None:
|
def save_model(self) -> None:
|
||||||
"""Save the model."""
|
"""Save the model."""
|
||||||
if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"):
|
if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2", "fsdpturbo"):
|
||||||
from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
||||||
|
|
||||||
DistributedPlugin(self.args.dist_config.name).save_model(
|
DistributedPlugin(self.args.dist_config.name).save_model(
|
||||||
|
|||||||
@@ -112,7 +112,6 @@ class ModelEngine:
|
|||||||
if self.args.custom_chat_template:
|
if self.args.custom_chat_template:
|
||||||
if not is_tokenizer(self.processor):
|
if not is_tokenizer(self.processor):
|
||||||
self.processor.chat_template = self.args.custom_chat_template
|
self.processor.chat_template = self.args.custom_chat_template
|
||||||
else:
|
|
||||||
tokenizer.chat_template = self.args.custom_chat_template
|
tokenizer.chat_template = self.args.custom_chat_template
|
||||||
|
|
||||||
def _init_model_config(self) -> HFConfig:
|
def _init_model_config(self) -> HFConfig:
|
||||||
@@ -184,7 +183,7 @@ class ModelEngine:
|
|||||||
if init_device.type == DeviceType.META:
|
if init_device.type == DeviceType.META:
|
||||||
assert self.args.quant_config is None, "Quantization is not supported with meta device."
|
assert self.args.quant_config is None, "Quantization is not supported with meta device."
|
||||||
with init_empty_weights():
|
with init_empty_weights():
|
||||||
model = AutoClass.from_config(self.model_config)
|
model = AutoClass.from_config(self.model_config, attn_implementation=self.args.flash_attn)
|
||||||
else:
|
else:
|
||||||
model = AutoClass.from_pretrained(
|
model = AutoClass.from_pretrained(
|
||||||
self.args.model,
|
self.args.model,
|
||||||
|
|||||||
@@ -250,7 +250,7 @@ class TrainingCheckpointCoordinator:
|
|||||||
num_training_steps=self._t.num_training_steps,
|
num_training_steps=self._t.num_training_steps,
|
||||||
)
|
)
|
||||||
|
|
||||||
if self._dist_name in ("fsdp2", "deepspeed"):
|
if self._dist_name in ("fsdp2", "fsdpturbo", "deepspeed"):
|
||||||
from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
||||||
|
|
||||||
DistributedPlugin(self._dist_name).save_checkpoint(
|
DistributedPlugin(self._dist_name).save_checkpoint(
|
||||||
@@ -306,7 +306,7 @@ class TrainingCheckpointCoordinator:
|
|||||||
self._t.global_step = metadata["global_step"]
|
self._t.global_step = metadata["global_step"]
|
||||||
self._t._resume_epoch = metadata["epoch"]
|
self._t._resume_epoch = metadata["epoch"]
|
||||||
|
|
||||||
if self._dist_name in ("fsdp2", "deepspeed"):
|
if self._dist_name in ("fsdp2", "fsdpturbo", "deepspeed"):
|
||||||
from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
|
||||||
|
|
||||||
DistributedPlugin(self._dist_name).load_checkpoint(
|
DistributedPlugin(self._dist_name).load_checkpoint(
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ from .base import KernelPlugin
|
|||||||
|
|
||||||
# Import built-in implementations so their class decorators populate the registry.
|
# Import built-in implementations so their class decorators populate the registry.
|
||||||
from .liger_kernel_ops import LigerKernel # noqa: F401
|
from .liger_kernel_ops import LigerKernel # noqa: F401
|
||||||
|
from .ops.linear_attention.fla import FlashLinearAttentionKernel # noqa: F401
|
||||||
from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel # noqa: F401
|
from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel # noqa: F401
|
||||||
from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel # noqa: F401
|
from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel # noqa: F401
|
||||||
from .ops.mlp.npu_swiglu import NpuSwiGluKernel # noqa: F401
|
from .ops.mlp.npu_swiglu import NpuSwiGluKernel # noqa: F401
|
||||||
|
|||||||
@@ -0,0 +1,103 @@
|
|||||||
|
# Copyright 2025 the LlamaFactory team.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Flash Linear Attention kernel plugin backed by FSDPTurbo's operator registry."""
|
||||||
|
|
||||||
|
from functools import partial
|
||||||
|
|
||||||
|
from ......accelerator.helper import DeviceType, get_current_accelerator
|
||||||
|
from ......utils import logging
|
||||||
|
from ......utils.types import HFModel
|
||||||
|
from ...base import BaseKernel, KernelPlugin
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
CHUNK_GATED_DELTA_RULE = "chunk_gated_delta_rule"
|
||||||
|
FUSED_RECURRENT_GATED_DELTA_RULE = "fused_recurrent_gated_delta_rule"
|
||||||
|
FLASH_LINEAR_ATTENTION_KERNELS = (
|
||||||
|
CHUNK_GATED_DELTA_RULE,
|
||||||
|
FUSED_RECURRENT_GATED_DELTA_RULE,
|
||||||
|
)
|
||||||
|
FLA_MODULE_ATTRIBUTES = {
|
||||||
|
CHUNK_GATED_DELTA_RULE: "chunk_gated_delta_rule",
|
||||||
|
FUSED_RECURRENT_GATED_DELTA_RULE: "recurrent_gated_delta_rule",
|
||||||
|
}
|
||||||
|
SUPPORTED_CHUNK_SIZES = (16, 32, 64)
|
||||||
|
|
||||||
|
|
||||||
|
@KernelPlugin("flash-linear-attention").register()
|
||||||
|
class FlashLinearAttentionKernel(BaseKernel):
|
||||||
|
"""Install selected FLA callables through FSDPTurbo's device operator registry."""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def check_device() -> None:
|
||||||
|
current = get_current_accelerator().type
|
||||||
|
if current not in (DeviceType.CUDA, DeviceType.NPU):
|
||||||
|
raise RuntimeError(f"FlashLinearAttentionKernel requires CUDA or NPU, current accelerator is {current}.")
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def check_deps() -> None:
|
||||||
|
try:
|
||||||
|
import fla.ops.gated_delta_rule # noqa: F401
|
||||||
|
import fsdp_turbo.ops.fla # noqa: F401
|
||||||
|
from fsdp_turbo.ops.registry import get_op # noqa: F401
|
||||||
|
from fsdp_turbo.utils.patch import patch_model_members # noqa: F401
|
||||||
|
except ImportError as exc:
|
||||||
|
raise RuntimeError("Flash Linear Attention and FSDPTurbo are required for this kernel.") from exc
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _apply(**kwargs) -> HFModel:
|
||||||
|
model = kwargs["model"]
|
||||||
|
config = kwargs.get("config") or {}
|
||||||
|
include_kernels = config.get("include_kernels", "auto")
|
||||||
|
chunk_size = config.get("chunk_size", 64)
|
||||||
|
|
||||||
|
if include_kernels == "auto" or include_kernels is True:
|
||||||
|
selected = list(FLASH_LINEAR_ATTENTION_KERNELS)
|
||||||
|
elif isinstance(include_kernels, str):
|
||||||
|
selected = [name.strip() for name in include_kernels.split(",") if name.strip()]
|
||||||
|
else:
|
||||||
|
raise TypeError("kernel_config.include_kernels must be 'auto' or a comma-separated string.")
|
||||||
|
|
||||||
|
if not selected:
|
||||||
|
raise ValueError("kernel_config.include_kernels must select at least one FLA kernel.")
|
||||||
|
|
||||||
|
unsupported = set(selected).difference(FLASH_LINEAR_ATTENTION_KERNELS)
|
||||||
|
if unsupported:
|
||||||
|
raise ValueError(f"Unsupported Flash Linear Attention kernels: {sorted(unsupported)}")
|
||||||
|
if isinstance(chunk_size, bool) or not isinstance(chunk_size, int) or chunk_size not in SUPPORTED_CHUNK_SIZES:
|
||||||
|
raise ValueError(f"chunk_size must be one of {SUPPORTED_CHUNK_SIZES}, got {chunk_size!r}.")
|
||||||
|
|
||||||
|
from fsdp_turbo.ops.registry import get_op
|
||||||
|
from fsdp_turbo.utils.patch import patch_model_members
|
||||||
|
|
||||||
|
patched = 0
|
||||||
|
named_modules = tuple(model.named_modules())
|
||||||
|
for op_name in selected:
|
||||||
|
module_attribute = FLA_MODULE_ATTRIBUTES[op_name]
|
||||||
|
op = get_op(op_name)
|
||||||
|
configured_op = partial(op, chunk_size=chunk_size) if op_name == CHUNK_GATED_DELTA_RULE else op
|
||||||
|
targets = {
|
||||||
|
f"{type(module).__module__}.{type(module).__name__}.{module_attribute}"
|
||||||
|
for _, module in named_modules
|
||||||
|
if callable(getattr(module, module_attribute, None))
|
||||||
|
}
|
||||||
|
matched = patch_model_members(model, sorted(targets), configured_op) if targets else 0
|
||||||
|
if matched == 0:
|
||||||
|
raise RuntimeError(f"FLA operator `{op_name}` did not match any model module attributes.")
|
||||||
|
patched += matched
|
||||||
|
|
||||||
|
logger.info_rank0(f"Flash Linear Attention kernels updated {patched} module callables: {selected}.")
|
||||||
|
return model
|
||||||
@@ -52,6 +52,14 @@ def get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int:
|
|||||||
return dist.get_rank(group) if group else 0
|
return dist.get_rank(group) if group else 0
|
||||||
|
|
||||||
|
|
||||||
|
def _get_text_position_ids(position_ids: Optional[Tensor]) -> Optional[Tensor]:
|
||||||
|
# Transformers < 5.4 broadcasts Qwen3.5 text positions over the mRoPE axes.
|
||||||
|
if position_ids is not None and position_ids.ndim == 3 and position_ids.stride(0) == 0:
|
||||||
|
position_ids = position_ids[0]
|
||||||
|
|
||||||
|
return position_ids.contiguous() if position_ids is not None and position_ids.ndim == 2 else None
|
||||||
|
|
||||||
|
|
||||||
class UlyssesAttention(torch.nn.Module):
|
class UlyssesAttention(torch.nn.Module):
|
||||||
"""Initialization.
|
"""Initialization.
|
||||||
|
|
||||||
@@ -123,8 +131,8 @@ class UlyssesAttention(torch.nn.Module):
|
|||||||
softmax_scale = q.shape[-1] ** -0.5
|
softmax_scale = q.shape[-1] ** -0.5
|
||||||
|
|
||||||
sp_world_size = get_ulysses_sequence_parallel_world_size(self.spg)
|
sp_world_size = get_ulysses_sequence_parallel_world_size(self.spg)
|
||||||
local_position_ids = position_ids
|
# HF FlashAttention only uses 2-D position IDs to detect packed sequences.
|
||||||
|
position_ids = _get_text_position_ids(position_ids)
|
||||||
if position_ids is not None:
|
if position_ids is not None:
|
||||||
global_position_ids = [torch.empty_like(position_ids) for _ in range(sp_world_size)]
|
global_position_ids = [torch.empty_like(position_ids) for _ in range(sp_world_size)]
|
||||||
dist.all_gather(global_position_ids, position_ids, group=self.spg)
|
dist.all_gather(global_position_ids, position_ids, group=self.spg)
|
||||||
@@ -144,13 +152,11 @@ class UlyssesAttention(torch.nn.Module):
|
|||||||
# contribute an all-ones shard.
|
# contribute an all-ones shard.
|
||||||
if torch.any(torch.stack(global_has_attention_mask)):
|
if torch.any(torch.stack(global_has_attention_mask)):
|
||||||
if attention_mask is None:
|
if attention_mask is None:
|
||||||
if local_position_ids is not None:
|
|
||||||
attention_mask = torch.ones_like(local_position_ids, dtype=torch.int64)
|
|
||||||
else:
|
|
||||||
attention_mask = torch.ones(query.shape[0], query.shape[1], dtype=torch.int64, device=query.device)
|
attention_mask = torch.ones(query.shape[0], query.shape[1], dtype=torch.int64, device=query.device)
|
||||||
else:
|
else:
|
||||||
attention_mask = attention_mask.to(torch.int64)
|
attention_mask = attention_mask.to(torch.int64)
|
||||||
|
|
||||||
|
attention_mask = attention_mask.contiguous()
|
||||||
global_attention_mask = [torch.empty_like(attention_mask) for _ in range(sp_world_size)]
|
global_attention_mask = [torch.empty_like(attention_mask) for _ in range(sp_world_size)]
|
||||||
dist.all_gather(global_attention_mask, attention_mask, group=self.spg)
|
dist.all_gather(global_attention_mask, attention_mask, group=self.spg)
|
||||||
attention_mask = torch.cat(global_attention_mask, dim=1).contiguous()
|
attention_mask = torch.cat(global_attention_mask, dim=1).contiguous()
|
||||||
|
|||||||
@@ -220,7 +220,7 @@ class FSDP2Engine:
|
|||||||
def is_lora_module_wrap(self, model) -> bool:
|
def is_lora_module_wrap(self, model) -> bool:
|
||||||
return any(isinstance(module, LoraLayer) for module in model.modules())
|
return any(isinstance(module, LoraLayer) for module in model.modules())
|
||||||
|
|
||||||
def prepare_model(self, model: HFModel) -> HFModel:
|
def prepare_model(self, model: HFModel, ignored_params: set[nn.Parameter] | None = None) -> HFModel:
|
||||||
if self.fsdp_mesh is None:
|
if self.fsdp_mesh is None:
|
||||||
logger.warning("No FSDP Mesh available, skipping FSDP wrapping.")
|
logger.warning("No FSDP Mesh available, skipping FSDP wrapping.")
|
||||||
return model
|
return model
|
||||||
@@ -236,6 +236,11 @@ class FSDP2Engine:
|
|||||||
names = ", ".join(cls.__name__ for cls in transformer_layer_cls_to_wrap)
|
names = ", ".join(cls.__name__ for cls in transformer_layer_cls_to_wrap)
|
||||||
logger.info(f"Applying per-layer FSDP to: {names}")
|
logger.info(f"Applying per-layer FSDP to: {names}")
|
||||||
|
|
||||||
|
def _ignored_params_for(module: nn.Module) -> set[nn.Parameter] | None:
|
||||||
|
if not ignored_params:
|
||||||
|
return None
|
||||||
|
return ignored_params.intersection(module.parameters()) or None
|
||||||
|
|
||||||
if self.is_lora_module_wrap(model):
|
if self.is_lora_module_wrap(model):
|
||||||
lora_modules = []
|
lora_modules = []
|
||||||
for module in model.modules():
|
for module in model.modules():
|
||||||
@@ -251,6 +256,7 @@ class FSDP2Engine:
|
|||||||
reshard_after_forward=self.reshard_after_forward,
|
reshard_after_forward=self.reshard_after_forward,
|
||||||
mp_policy=mp_policy,
|
mp_policy=mp_policy,
|
||||||
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
|
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
|
||||||
|
ignored_params=_ignored_params_for(module),
|
||||||
)
|
)
|
||||||
|
|
||||||
logger.info("Applying FSDP wrap for LoRA layer separately.")
|
logger.info("Applying FSDP wrap for LoRA layer separately.")
|
||||||
@@ -271,6 +277,7 @@ class FSDP2Engine:
|
|||||||
reshard_after_forward=self.reshard_after_forward,
|
reshard_after_forward=self.reshard_after_forward,
|
||||||
mp_policy=mp_policy,
|
mp_policy=mp_policy,
|
||||||
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
|
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
|
||||||
|
ignored_params=_ignored_params_for(module),
|
||||||
)
|
)
|
||||||
|
|
||||||
# BaseTrainer is the single source of truth for gradient checkpointing.
|
# BaseTrainer is the single source of truth for gradient checkpointing.
|
||||||
@@ -299,6 +306,7 @@ class FSDP2Engine:
|
|||||||
reshard_after_forward=self.reshard_after_forward,
|
reshard_after_forward=self.reshard_after_forward,
|
||||||
mp_policy=mp_policy,
|
mp_policy=mp_policy,
|
||||||
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
|
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
|
||||||
|
ignored_params=_ignored_params_for(model),
|
||||||
)
|
)
|
||||||
|
|
||||||
return model
|
return model
|
||||||
|
|||||||
@@ -0,0 +1,444 @@
|
|||||||
|
# Copyright 2025 the LlamaFactory team.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
|
||||||
|
|
||||||
|
from ....accelerator.interface import Dim, DistributedInterface
|
||||||
|
from ....utils.logging import get_logger
|
||||||
|
from ....utils.types import HFModel
|
||||||
|
from .fsdp2 import FSDP2Engine
|
||||||
|
|
||||||
|
|
||||||
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class FSDPTurboParallelState:
|
||||||
|
"""Own FSDPTurbo's expert topology without extending LlamaFactory's global interface."""
|
||||||
|
|
||||||
|
EDP = "edp"
|
||||||
|
EFSDP = "efsdp"
|
||||||
|
EP = "ep"
|
||||||
|
EXPERT_CP = "expert_cp"
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._initialized = False
|
||||||
|
self.dp_size = 1
|
||||||
|
self.cp_size = 1
|
||||||
|
self.ep_size = 1
|
||||||
|
self.efsdp_size = 1
|
||||||
|
self.edp_size = 1
|
||||||
|
self.expert_mesh: DeviceMesh | None = None
|
||||||
|
self.edp_mesh: DeviceMesh | None = None
|
||||||
|
self.efsdp_mesh: DeviceMesh | None = None
|
||||||
|
self.ep_mesh: DeviceMesh | None = None
|
||||||
|
self.expert_cp_mesh: DeviceMesh | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def initialized(self) -> bool:
|
||||||
|
return self._initialized
|
||||||
|
|
||||||
|
def initialize(self, dist_interface: DistributedInterface, dist_config: dict) -> None:
|
||||||
|
dp_size = dist_interface.get_world_size(Dim.DP)
|
||||||
|
cp_size = dist_interface.strategy.cp_size
|
||||||
|
ep_size = int(dist_config.get("ep_size", 1))
|
||||||
|
|
||||||
|
if ep_size < 1:
|
||||||
|
raise ValueError(f"ep_size must be positive, got {ep_size}.")
|
||||||
|
if dp_size % ep_size != 0:
|
||||||
|
raise ValueError(f"dp_size must be divisible by ep_size, got {dp_size} % {ep_size} != 0.")
|
||||||
|
|
||||||
|
topology = (dp_size, cp_size, ep_size)
|
||||||
|
if self._initialized:
|
||||||
|
current_topology = (self.dp_size, self.cp_size, self.ep_size)
|
||||||
|
if topology != current_topology:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"FSDPTurbo parallel state is already initialized with {current_topology}, got {topology}."
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
self.dp_size = dp_size
|
||||||
|
self.cp_size = cp_size
|
||||||
|
self.ep_size = ep_size
|
||||||
|
|
||||||
|
if ep_size > 1:
|
||||||
|
self.efsdp_size = dp_size // ep_size
|
||||||
|
self.edp_size = dp_size // (ep_size * self.efsdp_size)
|
||||||
|
if dist_interface.get_device_mesh(Dim.DP) is None:
|
||||||
|
raise RuntimeError("FSDPTurbo expert parallelism requires an initialized distributed device mesh.")
|
||||||
|
|
||||||
|
self.expert_mesh = init_device_mesh(
|
||||||
|
device_type=dist_interface.current_device.type,
|
||||||
|
mesh_shape=(self.edp_size, self.efsdp_size, self.ep_size, self.cp_size),
|
||||||
|
mesh_dim_names=(self.EDP, self.EFSDP, self.EP, self.EXPERT_CP),
|
||||||
|
)
|
||||||
|
self.edp_mesh = self.expert_mesh[self.EDP]
|
||||||
|
self.efsdp_mesh = self.expert_mesh[self.EFSDP]
|
||||||
|
self.ep_mesh = self.expert_mesh[self.EP]
|
||||||
|
self.expert_cp_mesh = self.expert_mesh[self.EXPERT_CP]
|
||||||
|
|
||||||
|
self._initialized = True
|
||||||
|
|
||||||
|
|
||||||
|
_FSDPTURBO_PARALLEL_STATE = FSDPTurboParallelState()
|
||||||
|
|
||||||
|
|
||||||
|
def get_fsdpturbo_parallel_state() -> FSDPTurboParallelState:
|
||||||
|
return _FSDPTURBO_PARALLEL_STATE
|
||||||
|
|
||||||
|
|
||||||
|
def _grad_to_local_fp32(grad: torch.Tensor) -> torch.Tensor:
|
||||||
|
from torch.distributed._tensor import DTensor
|
||||||
|
|
||||||
|
local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
|
||||||
|
return local_grad.detach().to(torch.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def _local_pth_sum(parameters: list[torch.nn.Parameter], norm_type: float, device: torch.device) -> torch.Tensor:
|
||||||
|
total = torch.zeros((), device=device, dtype=torch.float32)
|
||||||
|
for param in parameters:
|
||||||
|
grad = getattr(param, "grad", None)
|
||||||
|
if grad is None:
|
||||||
|
continue
|
||||||
|
total = total + torch.norm(_grad_to_local_fp32(grad), p=norm_type).pow(norm_type)
|
||||||
|
return total
|
||||||
|
|
||||||
|
|
||||||
|
def _allreduce_sum_(value: torch.Tensor, groups: list[object]) -> torch.Tensor:
|
||||||
|
import torch.distributed as dist
|
||||||
|
|
||||||
|
for group in groups:
|
||||||
|
if group is not None:
|
||||||
|
dist.all_reduce(value, op=dist.ReduceOp.SUM, group=group)
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
def clip_grad_norm_(model: HFModel, max_norm: float, **kwargs) -> float:
|
||||||
|
"""CP-aware grad norm clipping for FSDPTurbo EP + EFSDP + outer FSDP2.
|
||||||
|
|
||||||
|
Avoids torch.nn.utils.get_total_norm() since mixed DTensor meshes
|
||||||
|
(`dp` vs `efsdp`/`ep`) may hit DTensor stack propagation failures.
|
||||||
|
"""
|
||||||
|
from torch.distributed._tensor import DTensor
|
||||||
|
|
||||||
|
norm_type = float(kwargs.get("norm_type", 2.0))
|
||||||
|
dist_interface = DistributedInterface()
|
||||||
|
parallel_state = get_fsdpturbo_parallel_state()
|
||||||
|
if not parallel_state.initialized:
|
||||||
|
raise RuntimeError("FSDPTurbo parallel state must be initialized before clipping gradients.")
|
||||||
|
|
||||||
|
device = dist_interface.current_device
|
||||||
|
dp_group = dist_interface.get_group(Dim.DP)
|
||||||
|
cp_group = dist_interface.get_group(Dim.CP) if dist_interface.strategy.cp_size > 1 else None
|
||||||
|
ep_group = parallel_state.ep_mesh.get_group() if parallel_state.ep_mesh is not None else None
|
||||||
|
efsdp_group = parallel_state.efsdp_mesh.get_group() if parallel_state.efsdp_mesh is not None else None
|
||||||
|
expert_cp_group = (
|
||||||
|
parallel_state.expert_cp_mesh.get_group()
|
||||||
|
if parallel_state.expert_cp_mesh is not None and parallel_state.cp_size > 1
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
|
ep_params: list[torch.nn.Parameter] = []
|
||||||
|
non_ep_params: list[torch.nn.Parameter] = []
|
||||||
|
for param in model.parameters():
|
||||||
|
grad = getattr(param, "grad", None)
|
||||||
|
if grad is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
mesh_names = set(getattr(getattr(grad, "device_mesh", None), "mesh_dim_names", ()) or ())
|
||||||
|
is_ep_side = isinstance(grad, DTensor) and bool(mesh_names & {parallel_state.EP, parallel_state.EFSDP})
|
||||||
|
if is_ep_side:
|
||||||
|
ep_params.append(param)
|
||||||
|
else:
|
||||||
|
non_ep_params.append(param)
|
||||||
|
|
||||||
|
if not ep_params and not non_ep_params:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
total_pth = torch.zeros((), device=device, dtype=torch.float32)
|
||||||
|
if non_ep_params:
|
||||||
|
non_ep_pth = _local_pth_sum(non_ep_params, norm_type, device)
|
||||||
|
total_pth = total_pth + _allreduce_sum_(non_ep_pth, [dp_group, cp_group])
|
||||||
|
if ep_params:
|
||||||
|
ep_pth = _local_pth_sum(ep_params, norm_type, device)
|
||||||
|
total_pth = total_pth + _allreduce_sum_(ep_pth, [efsdp_group, ep_group, expert_cp_group])
|
||||||
|
|
||||||
|
total_norm = total_pth.pow(1.0 / norm_type)
|
||||||
|
clip_coef = min(max_norm / (float(total_norm.item()) + 1e-6), 1.0)
|
||||||
|
if clip_coef < 1.0:
|
||||||
|
for param in ep_params + non_ep_params:
|
||||||
|
grad = getattr(param, "grad", None)
|
||||||
|
if grad is not None:
|
||||||
|
grad.detach().mul_(clip_coef)
|
||||||
|
|
||||||
|
return float(total_norm.item())
|
||||||
|
|
||||||
|
|
||||||
|
def _get_model_type(model: HFModel) -> str | None:
|
||||||
|
return getattr(getattr(model, "config", None), "model_type", None)
|
||||||
|
|
||||||
|
|
||||||
|
class FSDPTurboEPModelSpec:
|
||||||
|
_registry: dict[str, "FSDPTurboEPModelSpec"] = {}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
ep_modules: list[str],
|
||||||
|
ep_fsdp_modules: list[str] | None = None,
|
||||||
|
prepare_fn: Callable[[HFModel], HFModel] | None = None,
|
||||||
|
) -> None:
|
||||||
|
self.ep_modules = ep_modules
|
||||||
|
self.ep_fsdp_modules = ep_fsdp_modules
|
||||||
|
self.prepare_fn = prepare_fn
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def register(
|
||||||
|
cls,
|
||||||
|
model_type: str,
|
||||||
|
ep_modules: list[str],
|
||||||
|
ep_fsdp_modules: list[str] | None = None,
|
||||||
|
):
|
||||||
|
def decorator(fn):
|
||||||
|
cls._registry[model_type] = cls(
|
||||||
|
ep_modules=ep_modules,
|
||||||
|
ep_fsdp_modules=ep_fsdp_modules,
|
||||||
|
prepare_fn=fn,
|
||||||
|
)
|
||||||
|
return fn
|
||||||
|
|
||||||
|
return decorator
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def get(cls, model: HFModel) -> "FSDPTurboEPModelSpec | None":
|
||||||
|
model_type = _get_model_type(model)
|
||||||
|
if model_type is None:
|
||||||
|
return None
|
||||||
|
return cls._registry.get(model_type)
|
||||||
|
|
||||||
|
def prepare(self, model: HFModel) -> HFModel:
|
||||||
|
if self.prepare_fn is None:
|
||||||
|
return model
|
||||||
|
return self.prepare_fn(model)
|
||||||
|
|
||||||
|
|
||||||
|
@FSDPTurboEPModelSpec.register(
|
||||||
|
"qwen3_moe",
|
||||||
|
ep_modules=["model.layers.{*}.mlp.experts"],
|
||||||
|
ep_fsdp_modules=["model.layers.{*}.mlp"],
|
||||||
|
)
|
||||||
|
def _prepare_qwen3_moe_for_ep(model: HFModel) -> HFModel:
|
||||||
|
prepared = 0
|
||||||
|
for module in model.modules():
|
||||||
|
if not all(hasattr(module, attr) for attr in ("gate_up_proj", "down_proj", "hidden_dim", "num_experts")):
|
||||||
|
continue
|
||||||
|
|
||||||
|
# FSDPTurbo's eager EP dispatcher expects sparse expert blocks to expose `hidden_size`.
|
||||||
|
if not hasattr(module, "hidden_size"):
|
||||||
|
module.hidden_size = module.hidden_dim
|
||||||
|
prepared += 1
|
||||||
|
|
||||||
|
if prepared:
|
||||||
|
logger.info_rank0(f"FSDPTurbo EP adapter: prepared {prepared} sparse expert modules for Transformers 5.x.")
|
||||||
|
else:
|
||||||
|
logger.info_rank0("FSDPTurbo EP adapter did not find a sparse expert module requiring preparation.")
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
@FSDPTurboEPModelSpec.register(
|
||||||
|
"qwen3_5_moe",
|
||||||
|
ep_modules=["model.language_model.layers.{*}.mlp.experts"],
|
||||||
|
ep_fsdp_modules=["model.language_model.layers.{*}.mlp"],
|
||||||
|
)
|
||||||
|
def _prepare_qwen3_5_moe_for_ep(model: HFModel) -> HFModel:
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
class FSDPTurboFSDP2Engine(FSDP2Engine):
|
||||||
|
"""FSDPTurbo EP adapter that reuses LlamaFactory's init/load flow.
|
||||||
|
|
||||||
|
Design:
|
||||||
|
- FSDPTurbo owns EP / EFSDP only.
|
||||||
|
- LlamaFactory owns FSDP / CP / init-load lifecycle.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, dist_config: dict, bf16: bool = False):
|
||||||
|
self.dist_config = dist_config
|
||||||
|
super().__init__(dist_config, bf16=bf16)
|
||||||
|
self.parallel_state = get_fsdpturbo_parallel_state()
|
||||||
|
self.parallel_state.initialize(self.dist_interface, self.dist_config)
|
||||||
|
self.ep_size = self.parallel_state.ep_size
|
||||||
|
self.ep_fsdp_size = self.parallel_state.efsdp_size
|
||||||
|
dp_mesh = self.dist_interface.get_device_mesh(Dim.DP)
|
||||||
|
if dp_mesh is not None:
|
||||||
|
self.fsdp_mesh = dp_mesh
|
||||||
|
logger.info(f"Using DP-orthogonal FSDP mesh: {self.fsdp_mesh}")
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _get_ep_fsdp_modules(spec: FSDPTurboEPModelSpec) -> list[str]:
|
||||||
|
if spec.ep_fsdp_modules is not None:
|
||||||
|
return spec.ep_fsdp_modules
|
||||||
|
|
||||||
|
ep_fsdp_modules = []
|
||||||
|
for module in spec.ep_modules:
|
||||||
|
if module.endswith(".experts"):
|
||||||
|
ep_fsdp_modules.append(module.removesuffix(".experts"))
|
||||||
|
else:
|
||||||
|
ep_fsdp_modules.append(module)
|
||||||
|
return ep_fsdp_modules
|
||||||
|
|
||||||
|
def shard_model(self, model: HFModel) -> HFModel:
|
||||||
|
"""Set storage dtype before FSDP materialization without leaking backend config into ModelEngine."""
|
||||||
|
param_dtype = torch.bfloat16 if self.mixed_precision == "bf16" else torch.float32
|
||||||
|
model = model.to(param_dtype)
|
||||||
|
logger.info_rank0(f"Using {param_dtype} for FSDPTurbo full tuning.")
|
||||||
|
return super().shard_model(model)
|
||||||
|
|
||||||
|
def _copy_weights(self, param, loaded_tensor):
|
||||||
|
"""Copy full checkpoint tensors into mixed-mesh DTensors from the inherited loader."""
|
||||||
|
from torch.distributed._tensor import DTensor, Shard
|
||||||
|
|
||||||
|
if loaded_tensor.dtype != param.dtype:
|
||||||
|
loaded_tensor = loaded_tensor.to(param.dtype)
|
||||||
|
|
||||||
|
if isinstance(param, DTensor):
|
||||||
|
local_tensor = param.to_local()
|
||||||
|
shard_placements = [
|
||||||
|
(i, placement) for i, placement in enumerate(param.placements) if isinstance(placement, Shard)
|
||||||
|
]
|
||||||
|
|
||||||
|
if not shard_placements:
|
||||||
|
local_tensor.copy_(loaded_tensor)
|
||||||
|
return
|
||||||
|
|
||||||
|
mesh = param.device_mesh
|
||||||
|
my_coordinate = mesh.get_coordinate()
|
||||||
|
if my_coordinate is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
sliced_tensor = loaded_tensor
|
||||||
|
for mesh_dim, shard_placement in shard_placements:
|
||||||
|
dim = shard_placement.dim
|
||||||
|
rank_in_dim = my_coordinate[mesh_dim]
|
||||||
|
world_size_in_dim = mesh.size(mesh_dim)
|
||||||
|
full_size = sliced_tensor.shape[dim]
|
||||||
|
chunk_size = (full_size + world_size_in_dim - 1) // world_size_in_dim
|
||||||
|
start = rank_in_dim * chunk_size
|
||||||
|
end = min(start + chunk_size, full_size)
|
||||||
|
|
||||||
|
if start >= full_size:
|
||||||
|
return
|
||||||
|
|
||||||
|
sliced_tensor = sliced_tensor.narrow(dim, start, end - start)
|
||||||
|
|
||||||
|
slices = [slice(None)] * local_tensor.ndim
|
||||||
|
for _, shard_placement in shard_placements:
|
||||||
|
dim = shard_placement.dim
|
||||||
|
slices[dim] = slice(0, sliced_tensor.shape[dim])
|
||||||
|
local_tensor[tuple(slices)].copy_(sliced_tensor)
|
||||||
|
return
|
||||||
|
|
||||||
|
param.data.copy_(loaded_tensor)
|
||||||
|
|
||||||
|
def prepare_model_ep(self, model: HFModel) -> tuple[HFModel, set]:
|
||||||
|
"""Apply FSDPTurbo EP/EFSDP and return parameters excluded from outer FSDP."""
|
||||||
|
from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import (
|
||||||
|
expert_fully_shard_modules,
|
||||||
|
)
|
||||||
|
from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules
|
||||||
|
from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig
|
||||||
|
from fsdp_turbo.utils.str_match import module_name_match
|
||||||
|
|
||||||
|
spec = FSDPTurboEPModelSpec.get(model)
|
||||||
|
if spec is None:
|
||||||
|
raise ValueError(f"No FSDPTurbo EP spec is registered for model_type={_get_model_type(model)}.")
|
||||||
|
|
||||||
|
ep_modules = spec.ep_modules
|
||||||
|
model = spec.prepare(model)
|
||||||
|
|
||||||
|
if self.ep_size > 1:
|
||||||
|
ep_plan = EPPlanConfig(
|
||||||
|
apply_modules=ep_modules,
|
||||||
|
dispatcher=self.dist_config.get("ep_dispatcher", "eager"),
|
||||||
|
apply_efsdp_modules=self._get_ep_fsdp_modules(spec),
|
||||||
|
)
|
||||||
|
ep_plan.gradient_divide_factor = float(self.ep_size * self.parallel_state.efsdp_size)
|
||||||
|
fsdp_plan = FSDPPlanConfig(
|
||||||
|
# FSDPTurbo uses this plan only to place EFSDP hooks and select its
|
||||||
|
# implementation. EFSDP targets come from ep_plan.apply_efsdp_modules.
|
||||||
|
apply_modules={},
|
||||||
|
hook_modules=self.dist_config.get("hook_modules", []),
|
||||||
|
fsdp_implementation=self.dist_config.get("fsdp_implementation", "native"),
|
||||||
|
)
|
||||||
|
ep_mesh = self.parallel_state.ep_mesh
|
||||||
|
efsdp_mesh = self.parallel_state.efsdp_mesh
|
||||||
|
if ep_mesh is None:
|
||||||
|
raise RuntimeError("FSDPTurbo EP mesh is not initialized.")
|
||||||
|
if self.ep_fsdp_size > 1 and efsdp_mesh is None:
|
||||||
|
raise RuntimeError("FSDPTurbo EFSDP mesh is not initialized.")
|
||||||
|
if self.rank == 0:
|
||||||
|
logger.info("Applying FSDPTurbo EP backend.")
|
||||||
|
logger.info(f"FSDPTurbo EP apply patterns: {ep_modules}")
|
||||||
|
logger.info(f"FSDPTurbo EP device mesh: {ep_mesh}")
|
||||||
|
logger.info(f"FSDPTurbo EP gradient divide factor: {ep_plan.gradient_divide_factor}")
|
||||||
|
|
||||||
|
model = expert_parallelize_modules(model, ep_mesh, ep_plan)
|
||||||
|
|
||||||
|
if self.ep_fsdp_size > 1:
|
||||||
|
if self.rank == 0:
|
||||||
|
logger.info(f"FSDPTurbo EFSDP apply patterns: {ep_plan.apply_efsdp_modules}")
|
||||||
|
logger.info(f"FSDPTurbo EFSDP device mesh: {efsdp_mesh}")
|
||||||
|
model = expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)
|
||||||
|
|
||||||
|
# Collect ignored params for the outer FSDP wrap
|
||||||
|
fsdp_ignored_modules = list(self.dist_config.get("fsdp_ignored_modules", []))
|
||||||
|
if self.ep_size > 1:
|
||||||
|
fsdp_ignored_modules.extend(ep_modules)
|
||||||
|
|
||||||
|
ignored_params = set()
|
||||||
|
if fsdp_ignored_modules:
|
||||||
|
for name, module in model.named_modules():
|
||||||
|
for pattern in fsdp_ignored_modules:
|
||||||
|
if module_name_match(pattern, name):
|
||||||
|
ignored_params.update(list(module.parameters(recurse=True)))
|
||||||
|
|
||||||
|
if ignored_params and self.rank == 0:
|
||||||
|
logger.info(f"FSDPTurbo FSDP2: Ignoring {len(ignored_params)} EP parameters in outer FSDP.")
|
||||||
|
|
||||||
|
return model, ignored_params
|
||||||
|
|
||||||
|
def prepare_model(self, model: HFModel) -> HFModel:
|
||||||
|
# Apply FSDPTurbo EP first, then shard the remaining parameters with LlamaFactory FSDP2.
|
||||||
|
model, ignored_params = self.prepare_model_ep(model)
|
||||||
|
return super().prepare_model(model, ignored_params=ignored_params)
|
||||||
|
|
||||||
|
def _warmup_grad_norm(self, model: HFModel) -> None:
|
||||||
|
"""Warm up collectives without stacking gradients from different DTensor meshes."""
|
||||||
|
if self.fsdp_mesh is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info_rank0("Warming up FSDPTurbo mixed-mesh grad norm computation...")
|
||||||
|
for param in model.parameters():
|
||||||
|
if param.requires_grad:
|
||||||
|
param.grad = torch.zeros_like(param)
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
clip_grad_norm_(model, 1.0)
|
||||||
|
|
||||||
|
for param in model.parameters():
|
||||||
|
if param.requires_grad:
|
||||||
|
param.grad = None
|
||||||
|
|
||||||
|
logger.info_rank0("FSDPTurbo mixed-mesh grad norm warmup completed.")
|
||||||
@@ -20,7 +20,7 @@ reads mesh topology from ``TrainingArguments`` and never puts it in backend para
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from dataclasses import asdict, dataclass
|
from dataclasses import asdict, dataclass, field
|
||||||
from typing import TYPE_CHECKING, Literal
|
from typing import TYPE_CHECKING, Literal
|
||||||
|
|
||||||
from ....utils.plugin import BasePlugin
|
from ....utils.plugin import BasePlugin
|
||||||
@@ -41,6 +41,24 @@ class FSDP2Params:
|
|||||||
dcp_path: str | None = None
|
dcp_path: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class FSDPTurboParams:
|
||||||
|
name: Literal["fsdpturbo"] = "fsdpturbo"
|
||||||
|
reshard_after_forward: bool = True
|
||||||
|
offload_params: bool = False
|
||||||
|
pin_memory: bool = True
|
||||||
|
dcp_path: str | None = None
|
||||||
|
ep_size: int = 1
|
||||||
|
ep_dispatcher: str = "eager"
|
||||||
|
fsdp_ignored_modules: list[str] = field(default_factory=list)
|
||||||
|
hook_modules: list[str] = field(default_factory=list)
|
||||||
|
fsdp_implementation: str = "native"
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
if self.ep_size < 1:
|
||||||
|
raise ValueError(f"ep_size must be positive, got {self.ep_size}.")
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class DeepSpeedParams:
|
class DeepSpeedParams:
|
||||||
name: Literal["deepspeed"] = "deepspeed"
|
name: Literal["deepspeed"] = "deepspeed"
|
||||||
@@ -83,6 +101,40 @@ class FSDP2Distributed(BaseDistributed):
|
|||||||
load_checkpoint(model, optimizer, ckpt_dir, **kwargs)
|
load_checkpoint(model, optimizer, ckpt_dir, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
@DistributedPlugin("fsdpturbo").register()
|
||||||
|
class FSDPTurboDistributed(BaseDistributed):
|
||||||
|
@staticmethod
|
||||||
|
def shard_model(model: HFModel, dist_config: PluginConfig | FSDPTurboParams, **kwargs) -> HFModel:
|
||||||
|
dist_config = DistributedPlugin.parse_params(dist_config, FSDPTurboParams)
|
||||||
|
from .fsdpturbo import FSDPTurboFSDP2Engine
|
||||||
|
|
||||||
|
return FSDPTurboFSDP2Engine(asdict(dist_config), bf16=bool(kwargs.get("bf16"))).shard_model(model)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def clip_grad_norm(model: HFModel, max_norm: float, **kwargs) -> float:
|
||||||
|
from .fsdpturbo import clip_grad_norm_
|
||||||
|
|
||||||
|
return clip_grad_norm_(model, max_norm, **kwargs)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def save_model(model, output_dir, processor) -> None:
|
||||||
|
from .fsdp2 import save_model
|
||||||
|
|
||||||
|
save_model(model, output_dir, processor)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def save_checkpoint(model, optimizer, ckpt_dir, **kwargs) -> None:
|
||||||
|
from .fsdp2 import save_checkpoint
|
||||||
|
|
||||||
|
save_checkpoint(model, optimizer, ckpt_dir, **kwargs)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def load_checkpoint(model, optimizer, ckpt_dir, **kwargs) -> None:
|
||||||
|
from .fsdp2 import load_checkpoint
|
||||||
|
|
||||||
|
load_checkpoint(model, optimizer, ckpt_dir, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
@DistributedPlugin("deepspeed").register()
|
@DistributedPlugin("deepspeed").register()
|
||||||
class DeepSpeedDistributed(BaseDistributed):
|
class DeepSpeedDistributed(BaseDistributed):
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
|||||||
@@ -26,7 +26,9 @@ def test_get_args_from_yaml(tmp_path: Path):
|
|||||||
trust_remote_code: true
|
trust_remote_code: true
|
||||||
model_class: llm
|
model_class: llm
|
||||||
kernel_config:
|
kernel_config:
|
||||||
name: auto
|
name: auto, flash-linear-attention
|
||||||
|
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
|
||||||
|
chunk_size: 32
|
||||||
peft_config:
|
peft_config:
|
||||||
name: lora
|
name: lora
|
||||||
r: 8
|
r: 8
|
||||||
@@ -58,7 +60,11 @@ def test_get_args_from_yaml(tmp_path: Path):
|
|||||||
model_args, data_args, training_args, sample_args = get_args()
|
model_args, data_args, training_args, sample_args = get_args()
|
||||||
assert data_args.train_dataset == "llamafactory/v1-sft-demo"
|
assert data_args.train_dataset == "llamafactory/v1-sft-demo"
|
||||||
assert model_args.model == "llamafactory/tiny-random-qwen3"
|
assert model_args.model == "llamafactory/tiny-random-qwen3"
|
||||||
assert model_args.kernel_config.name == "auto"
|
assert model_args.kernel_config.name == "auto, flash-linear-attention"
|
||||||
|
assert model_args.kernel_config.get("include_kernels") == (
|
||||||
|
"chunk_gated_delta_rule, fused_recurrent_gated_delta_rule"
|
||||||
|
)
|
||||||
|
assert model_args.kernel_config.get("chunk_size") == 32
|
||||||
assert model_args.peft_config.name == "lora"
|
assert model_args.peft_config.name == "lora"
|
||||||
assert model_args.peft_config.get("r") == 8
|
assert model_args.peft_config.get("r") == 8
|
||||||
assert training_args.output_dir == "outputs/test_run"
|
assert training_args.output_dir == "outputs/test_run"
|
||||||
@@ -68,3 +74,16 @@ def test_get_args_from_yaml(tmp_path: Path):
|
|||||||
assert training_args.bf16 is False
|
assert training_args.bf16 is False
|
||||||
assert training_args.dist_config is None
|
assert training_args.dist_config is None
|
||||||
assert sample_args.sample_backend == "hf"
|
assert sample_args.sample_backend == "hf"
|
||||||
|
|
||||||
|
|
||||||
|
def test_qwen35_fsdpturbo_example_uses_v1_arguments():
|
||||||
|
config_file = (
|
||||||
|
Path(__file__).parents[2] / "examples" / "v1" / "train_full" / "train_full_qwen3_moe_fsdpturbo_ep_fsdp.yaml"
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch.object(sys, "argv", ["test_args_parser.py", str(config_file)]):
|
||||||
|
model_args, _, training_args, _ = get_args()
|
||||||
|
|
||||||
|
assert model_args.model == "Qwen/Qwen3.5-35B-A3B"
|
||||||
|
assert model_args.custom_chat_template is None
|
||||||
|
assert training_args.dist_config.name == "fsdpturbo"
|
||||||
|
|||||||
@@ -13,12 +13,32 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
import sys
|
import sys
|
||||||
|
from functools import partial
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
import torch.multiprocessing as mp
|
import torch.multiprocessing as mp
|
||||||
|
from torch import nn
|
||||||
from transformers import AutoModelForCausalLM
|
from transformers import AutoModelForCausalLM
|
||||||
|
|
||||||
|
|
||||||
|
def _original_fla_op(*args, **kwargs):
|
||||||
|
return args, kwargs
|
||||||
|
|
||||||
|
|
||||||
|
class _LinearAttention(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.chunk_gated_delta_rule = _original_fla_op
|
||||||
|
self.recurrent_gated_delta_rule = _original_fla_op
|
||||||
|
|
||||||
|
|
||||||
|
class _FLAModel(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.linear_attn = _LinearAttention()
|
||||||
|
|
||||||
|
|
||||||
def _apply_kernel(rank) -> None:
|
def _apply_kernel(rank) -> None:
|
||||||
with patch("torch.accelerator.current_accelerator") as mock_get_accelerator:
|
with patch("torch.accelerator.current_accelerator") as mock_get_accelerator:
|
||||||
mock_device = MagicMock()
|
mock_device = MagicMock()
|
||||||
@@ -73,3 +93,62 @@ def test_apply_kernel():
|
|||||||
|
|
||||||
def test_apply_all_kernels():
|
def test_apply_all_kernels():
|
||||||
mp.spawn(_apply_all_kernels)
|
mp.spawn(_apply_all_kernels)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.runs_on(["npu"])
|
||||||
|
def test_flash_linear_attention_kernels_compose_with_auto(monkeypatch):
|
||||||
|
import fsdp_turbo.ops.fla # noqa: F401
|
||||||
|
from fsdp_turbo.ops import get_op
|
||||||
|
|
||||||
|
from llamafactory.v1.plugins.model_plugins.kernels import interface
|
||||||
|
from llamafactory.v1.plugins.model_plugins.kernels.ops.linear_attention.fla import (
|
||||||
|
FlashLinearAttentionKernel,
|
||||||
|
)
|
||||||
|
|
||||||
|
model = _FLAModel()
|
||||||
|
auto_calls = []
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
interface,
|
||||||
|
"_apply_auto_kernels",
|
||||||
|
lambda model, **kwargs: auto_calls.append((model, kwargs)) or model,
|
||||||
|
)
|
||||||
|
# FLA execution is outside this bridge test; its external runtime is not required.
|
||||||
|
monkeypatch.setattr(FlashLinearAttentionKernel, "check_deps", staticmethod(lambda: None))
|
||||||
|
|
||||||
|
config = {
|
||||||
|
"name": "auto, flash-linear-attention",
|
||||||
|
"include_kernels": "fused_recurrent_gated_delta_rule, chunk_gated_delta_rule",
|
||||||
|
"chunk_size": 32,
|
||||||
|
}
|
||||||
|
assert interface.apply_kernels(model, config) is model
|
||||||
|
assert auto_calls == [(model, {"config": config, "require_logits": False})]
|
||||||
|
assert get_op("chunk_gated_delta_rule").__module__ == "fsdp_turbo.ops.fla"
|
||||||
|
|
||||||
|
chunk_op = model.linear_attn.chunk_gated_delta_rule
|
||||||
|
assert isinstance(chunk_op, partial)
|
||||||
|
assert chunk_op.func.__module__ == "fsdp_turbo.ops.fla"
|
||||||
|
assert chunk_op.keywords == {"chunk_size": 32}
|
||||||
|
assert model.linear_attn.recurrent_gated_delta_rule.__module__ == "fsdp_turbo.ops.fla"
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="did not match any model module attributes"):
|
||||||
|
FlashLinearAttentionKernel.apply(
|
||||||
|
model=nn.Linear(2, 2),
|
||||||
|
config={"include_kernels": "chunk_gated_delta_rule", "chunk_size": 32},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_flash_linear_attention_kernel_validates_config(monkeypatch):
|
||||||
|
from llamafactory.v1.plugins.model_plugins.kernels.ops.linear_attention.fla import (
|
||||||
|
FlashLinearAttentionKernel,
|
||||||
|
)
|
||||||
|
|
||||||
|
model = nn.Sequential(nn.Linear(2, 2))
|
||||||
|
monkeypatch.setattr(FlashLinearAttentionKernel, "check_device", staticmethod(lambda: None))
|
||||||
|
monkeypatch.setattr(FlashLinearAttentionKernel, "check_deps", staticmethod(lambda: None))
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="chunk_size"):
|
||||||
|
FlashLinearAttentionKernel.apply(model=model, config={"include_kernels": "auto", "chunk_size": 48})
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Unsupported Flash Linear Attention kernels"):
|
||||||
|
FlashLinearAttentionKernel.apply(model=model, config={"include_kernels": "not_a_kernel"})
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ from llamafactory.v1.accelerator.interface import DistributedInterface
|
|||||||
from llamafactory.v1.config.model_args import ModelArguments
|
from llamafactory.v1.config.model_args import ModelArguments
|
||||||
from llamafactory.v1.config.training_args import TrainingArguments
|
from llamafactory.v1.config.training_args import TrainingArguments
|
||||||
from llamafactory.v1.core.model_engine import ModelEngine
|
from llamafactory.v1.core.model_engine import ModelEngine
|
||||||
|
from llamafactory.v1.plugins.model_plugins.parallelization import ulysses
|
||||||
from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
|
from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
|
||||||
SequenceParallelModelPlugin,
|
SequenceParallelModelPlugin,
|
||||||
sequence_parallel_loss,
|
sequence_parallel_loss,
|
||||||
@@ -28,6 +29,39 @@ from llamafactory.v1.utils.env import find_available_port
|
|||||||
from llamafactory.v1.utils.pytest import dist_env
|
from llamafactory.v1.utils.pytest import dist_env
|
||||||
|
|
||||||
|
|
||||||
|
def test_qwen3_5_broadcast_position_ids_keep_packed_boundaries(monkeypatch: pytest.MonkeyPatch):
|
||||||
|
local_position_ids = torch.tensor([[0, 1, 0]])
|
||||||
|
remote_position_ids = torch.tensor([[1, 2, 3]])
|
||||||
|
mrope_position_ids = local_position_ids.unsqueeze(0).expand(3, -1, -1)
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", lambda _, tensor, *__: tensor)
|
||||||
|
monkeypatch.setattr(ulysses, "get_ulysses_sequence_parallel_world_size", lambda _: 2)
|
||||||
|
|
||||||
|
def fake_all_gather(outputs, tensor, **_):
|
||||||
|
outputs[0].copy_(tensor)
|
||||||
|
outputs[1].copy_(remote_position_ids if tensor.shape == local_position_ids.shape else tensor)
|
||||||
|
|
||||||
|
def fake_attention(query, _key, _value, _attention_mask, **kwargs):
|
||||||
|
captured["position_ids"] = kwargs["position_ids"]
|
||||||
|
return query
|
||||||
|
|
||||||
|
monkeypatch.setattr(ulysses.dist, "all_gather", fake_all_gather)
|
||||||
|
attention = ulysses.UlyssesAttention(sequence_process_group=object(), attn_fn=fake_attention)
|
||||||
|
hidden_states = torch.zeros(1, 3, 2, 4)
|
||||||
|
|
||||||
|
attention(hidden_states, hidden_states, hidden_states, None, 6, position_ids=mrope_position_ids)
|
||||||
|
|
||||||
|
assert captured["position_ids"].tolist() == [[0, 1, 0, 1, 2, 3]]
|
||||||
|
assert captured["position_ids"].is_contiguous()
|
||||||
|
|
||||||
|
|
||||||
|
def test_true_mrope_position_ids_are_not_used_as_packed_boundaries():
|
||||||
|
mrope_position_ids = torch.tensor([[[0, 1, 2]], [[0, 1, 1]], [[0, 1, 0]]])
|
||||||
|
|
||||||
|
assert ulysses._get_text_position_ids(mrope_position_ids) is None
|
||||||
|
|
||||||
|
|
||||||
def _test_sequence_parallel_loss(
|
def _test_sequence_parallel_loss(
|
||||||
local_rank: int, world_size: int, master_port: int, cp_size: int, dp_size: int, batch_size: int
|
local_rank: int, world_size: int, master_port: int, cp_size: int, dp_size: int, batch_size: int
|
||||||
):
|
):
|
||||||
|
|||||||
@@ -0,0 +1,134 @@
|
|||||||
|
# Copyright 2025 the LlamaFactory team.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from llamafactory.v1.plugins.trainer_plugins.distributed import fsdpturbo as fsdpturbo_module
|
||||||
|
from llamafactory.v1.plugins.trainer_plugins.distributed.fsdpturbo import (
|
||||||
|
FSDPTurboEPModelSpec,
|
||||||
|
FSDPTurboFSDP2Engine,
|
||||||
|
FSDPTurboParallelState,
|
||||||
|
)
|
||||||
|
from llamafactory.v1.plugins.trainer_plugins.distributed.interface import (
|
||||||
|
DistributedPlugin,
|
||||||
|
FSDPTurboParams,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _Model(torch.nn.Module):
|
||||||
|
def __init__(self, model_type: str):
|
||||||
|
super().__init__()
|
||||||
|
self.config = SimpleNamespace(model_type=model_type)
|
||||||
|
|
||||||
|
|
||||||
|
def test_qwen35_ep_model_spec():
|
||||||
|
spec = FSDPTurboEPModelSpec.get(_Model("qwen3_5_moe"))
|
||||||
|
|
||||||
|
assert spec is not None
|
||||||
|
assert spec.ep_modules == ["model.language_model.layers.{*}.mlp.experts"]
|
||||||
|
assert spec.ep_fsdp_modules == ["model.language_model.layers.{*}.mlp"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_fsdpturbo_uses_class_plugin_and_strict_backend_params():
|
||||||
|
plugin = DistributedPlugin("fsdpturbo")
|
||||||
|
params = plugin.parse_params({"name": "fsdpturbo", "ep_size": 4}, FSDPTurboParams)
|
||||||
|
|
||||||
|
assert params.ep_size == 4
|
||||||
|
assert callable(plugin.shard_model)
|
||||||
|
assert callable(plugin.clip_grad_norm)
|
||||||
|
with pytest.raises(ValueError, match="Unknown params"):
|
||||||
|
plugin.parse_params({"name": "fsdpturbo", "cp_size": 2}, FSDPTurboParams)
|
||||||
|
for key in ("ep_modules", "ep_fsdp_modules"):
|
||||||
|
with pytest.raises(ValueError, match="Unknown params"):
|
||||||
|
plugin.parse_params({"name": "fsdpturbo", key: ["model.layers.*.mlp"]}, FSDPTurboParams)
|
||||||
|
|
||||||
|
|
||||||
|
def test_fsdpturbo_sets_storage_dtype_inside_backend(monkeypatch):
|
||||||
|
from llamafactory.v1.plugins.trainer_plugins.distributed.fsdp2 import FSDP2Engine
|
||||||
|
|
||||||
|
monkeypatch.setattr(FSDP2Engine, "shard_model", lambda self, model: model)
|
||||||
|
engine = object.__new__(FSDPTurboFSDP2Engine)
|
||||||
|
engine.mixed_precision = "bf16"
|
||||||
|
model = torch.nn.Linear(2, 2, dtype=torch.float32)
|
||||||
|
|
||||||
|
assert engine.shard_model(model).weight.dtype == torch.bfloat16
|
||||||
|
|
||||||
|
|
||||||
|
def test_fsdpturbo_sets_public_efsdp_gradient_divide_factor(monkeypatch):
|
||||||
|
expert_parallel_module = pytest.importorskip("fsdp_turbo.distributed.expert_parallel.expert_parallel")
|
||||||
|
expert_fully_shard_module = pytest.importorskip(
|
||||||
|
"fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel"
|
||||||
|
)
|
||||||
|
captured = {}
|
||||||
|
monkeypatch.setattr(expert_parallel_module, "expert_parallelize_modules", lambda model, mesh, plan: model)
|
||||||
|
|
||||||
|
def _expert_fully_shard_modules(model, mesh, ep_plan, fsdp_plan):
|
||||||
|
captured["gradient_divide_factor"] = ep_plan.gradient_divide_factor
|
||||||
|
return model
|
||||||
|
|
||||||
|
monkeypatch.setattr(expert_fully_shard_module, "expert_fully_shard_modules", _expert_fully_shard_modules)
|
||||||
|
|
||||||
|
engine = object.__new__(FSDPTurboFSDP2Engine)
|
||||||
|
engine.dist_config = {"ep_dispatcher": "eager"}
|
||||||
|
engine.ep_size = 4
|
||||||
|
engine.ep_fsdp_size = 2
|
||||||
|
engine.parallel_state = SimpleNamespace(efsdp_size=2, ep_mesh=object(), efsdp_mesh=object())
|
||||||
|
engine.rank = 0
|
||||||
|
|
||||||
|
engine.prepare_model_ep(_Model("qwen3_5_moe"))
|
||||||
|
|
||||||
|
assert captured["gradient_divide_factor"] == 8.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_fsdpturbo_owns_expert_mesh_topology(monkeypatch):
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
class _Mesh:
|
||||||
|
def __init__(self, name="expert"):
|
||||||
|
self.name = name
|
||||||
|
|
||||||
|
def __getitem__(self, name):
|
||||||
|
return _Mesh(name)
|
||||||
|
|
||||||
|
def _init_device_mesh(**kwargs):
|
||||||
|
calls.append(kwargs)
|
||||||
|
return _Mesh()
|
||||||
|
|
||||||
|
class _DistributedInterface:
|
||||||
|
current_device = torch.device("cpu")
|
||||||
|
strategy = SimpleNamespace(cp_size=1)
|
||||||
|
|
||||||
|
def get_world_size(self, dim):
|
||||||
|
return 16
|
||||||
|
|
||||||
|
def get_device_mesh(self, dim):
|
||||||
|
return _Mesh("dp")
|
||||||
|
|
||||||
|
monkeypatch.setattr(fsdpturbo_module, "init_device_mesh", _init_device_mesh)
|
||||||
|
state = FSDPTurboParallelState()
|
||||||
|
state.initialize(_DistributedInterface(), {"ep_size": 8})
|
||||||
|
|
||||||
|
assert calls == [
|
||||||
|
{
|
||||||
|
"device_type": "cpu",
|
||||||
|
"mesh_shape": (1, 2, 8, 1),
|
||||||
|
"mesh_dim_names": ("edp", "efsdp", "ep", "expert_cp"),
|
||||||
|
}
|
||||||
|
]
|
||||||
|
assert state.ep_mesh.name == "ep"
|
||||||
|
assert state.efsdp_mesh.name == "efsdp"
|
||||||
|
assert state.expert_cp_mesh.name == "expert_cp"
|
||||||
Reference in New Issue
Block a user