| [refactor] breaking update ! refactor config system to adapt torchtitan main branch
Co-authored-by: 1Fire4<wangdingyi2@huawei.com>
# message auto-generated for no-merge-commit merge:
!175 merge refactor/adapt_dsv3 into master
[refactor] breaking update ! refactor config system to adapt torchtitan main branch
Created-by: hitwdy
Commit-by: 1Fire4
Merged-by: cann-robot
Description: ## 描述
### 🚨 [破坏性更新] 配置系统全面重构与模型架构展平
> **⚠️ 注意:这是一次破坏性更新,将导致目前尚未完成重构适配的模型完全无法运行。当前本 PR 仅完成了对 DeepSeek-V3 的适配,后续将尽快完成其余模型的适配工作。**
### 💡 重构前后的核心用法差异 (Before vs. After)
本次重构深度对齐了上游最新的基于 Python Dataclass 的 Config Registry 机制,彻底抛弃了过往僵化的 TOML 配置文件。在用户层面,最直观的使用体验变化如下:
| 维度 | 重构前 (Before) | 重构后 (After) |
|---|---|---|
| **配置入口** | 强依赖硬编码的 TOML 文件(如 train_configs/*.toml)以及庞大、扁平的全局 JobConfig 类。 | **完全废弃 TOML**。全面拥抱 Python 原生的 config_registry.py,各个组件自带 Config 类嵌套,具备严格的类型安全。 |
| **NPU 专属参数** | 依赖维护一个独立的 custom_config.py,并且需要通过运行时的 _merge_configs 函数把参数强行“挂载”到 JobConfig 中。 | **类继承式无缝注入**。直接继承上游的原生配置类(如 ParallelismConfig),NPU 专属特性(如 Ulysses CP、交换优化器等)作为原生 Field 直接定义,取消外部 Merge 操作。 |
| **命令行交互** | 命令行参数传递和覆盖依赖定制化的平铺 TOML 键值解析。 | **基于 tyro 的嵌套解析**。得益于 Dataclass,所有的 NPU 专属参数现在可以像原生参数一样,通过带层级的命令行参数直接覆盖(例如 --training.swap_optimizer)。 |
| **代码文件组织** | 并行配置与模型定义分散在不同的深层嵌套目录(model/,infra/),导致重度依赖和冗长的 import 路径。 | **目录全面扁平化**。model.py、parallelize.py、state_dict_adapter.py 等核心组件全部被上提至各个模型的根目录下(如 models/deepseek_v3/),架构一目了然。 |
---
### 🔧 五大核心适配与重构要点详解
#### 1. 配置系统 (Config System) 的深度融合与精简
这是最核心的底层改造之一,彻底改变了 NPU 专属参数的挂载与流转方式:
* **废弃 custom_config.py**:移除了脱离上游体系的独立 NPU 自定义配置逻辑。
* **配置类继承式覆盖**:在 torchtitan_npu/config/configs.py 中,直接采用继承上游配置类(如 ParallelismConfig, OptimizerConfig)的设计模式。
* **CLI 原生暴露**:将 NPU 的专属特性(如 enable_custom_context_parallel、swap_optimizer、match_rms_adamw 等)作为新增字段(Fields)直接注入到继承类中。通过 tyro 解析器,这使得 NPU 专属参数可以和原生参数一样完美兼容命令行修改。
#### 2. 模型目录结构的全面展平 (Directory Flattening)
对齐上游最新的模型目录规范,对 deepseek_v3 的文件树进行了“去嵌套化”处理:
* **消除冗余层级**:将原本深层嵌套的 model/model.py 直接迁移至模型根目录。
* **统一组件入口**:原有的并行配置代码 infra/parallelize.py 上提至模型根目录并直接命名为 parallelize.py;模型权重适配器 state_dict_adapter.py 同样上提至第一层级。
* **清理废弃配置**:删除各个模型目录下大量冗余的、用于测试的硬编码 TOML 文件(如 train_configs/deepseek_v32_671b_debug.toml 等),极大地净化了代码库。
#### 3. 模型内部逻辑与类的重构
* **拥抱 Dataclass Component**:废弃了原有单一且臃肿的 model_args 结构,修改当前的内部定义以对齐上游范式。为每一个可配置的模型组件(Component)引入独立的内部 Config 数据类,确保配置项流转的高内聚与低耦合。
#### 4. 补丁层 (Patches & Converters) 的自适应修正
由于配置文件的入口与加载路径发生了根本性的改变,底层的 AST 拦截与替换逻辑也进行了同步升级:
* **动态参数解包适配**:大幅更新 converters/framework/model_custom_config_converter.py 以及 quant_converter.py,确保其能平滑兼容全新的继承式 Config 结构。
* **上游核心补丁更新**:修改了 patches/torchtitan/ 目录下的核心文件(如 expert_parallel.py, hf_datasets.py, loss.py 等)。
* **新增状态劫持机制**:新增了 _trainer_config_stash.py 文件,用于在 NPU 训练生命周期内更加稳妥地暂存和劫持 Trainer 的运行期环境状态。
#### 5. 外围生态 (Tools, Scripts & Docs) 的全局翻新
* **入口与工具链更新**:对主入口 torchtitan_npu/train.py 以及所有周边诊断工具(checkpoint_patch.py, flight_recorder.py, profiling.py)中的 import 依赖路径进行了地毯式替换。
* **启动脚本规范化**:重构 scripts/run_train.sh 和 run_train_multinodes.sh,使其完美兼容最新的、基于 tyro 的命令行参数解析方式。
## 类型
- [ ] Bug 修复
- [ ] 新功能
- [x] 重构(即不是新增功能,也不是修改bug的代码变动)
- [x] 构建过程或辅助工具的变动
- [x] 文档内容更新
## Checklist:
- [x] 我的代码遵循这个项目的代码风格
- [x] 我已经自己测试过我的代码
- [x] 我已经更新了相应的文档
- [x] 我已经在标题中正确使用了类型标签(例如:feat, fix, refactor, docs, test)
## 如何测试
本PR依赖python3.11及特定版本torch,torchtitan,torch_npu
```
# 1. torchtitan
git clone https://github.com/pytorch/torchtitan.git
cd torchtitan
git checkout ac13e536c84e7f6647b14fa9375c3c8a8a2b8578
cd <torchtitan_npu>
ln -s <path>/torchtitan/torchtitan .
# 2.torch
pip3 install --no-cache-dir "https://download-r2.pytorch.org/whl/nightly/cpu/torch-2.12.0.dev20260317%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl"
# 3. torch_npu
git clone https://gitcode.com/Ascend/pytorch.git /tmp/pytorch_npu && git -C /tmp/pytorch_npu checkout f9cbf1f179b59e75b915a72cfc3187f0aadfdea3
cd /tmp/pytorch_npu && bash ci/build.sh --python=3.11
pip3 install --upgrade /tmp/pytorch_npu/dist/torch_npu*.whl
```
## 其他信息
剩余待办:
[RFC](https://gitcode.com/cann/torchtitan-npu/issues/21)
See merge request: cann/torchtitan-npu!175 | 18 天前 |