Repository navigation
fix: preserve dtype across pipeline parallel communication - #255
JYMiracle305 wants to merge 2 commits into
Conversation
| void ISend::SetupContext( | ||
| const std::vector<std::shared_ptr<Tensor>> &input_tensors, | ||
| const std::vector<std::shared_ptr<Tensor>> &output_tensors) { | ||
| CHECK_EQ(input_tensors.size(), shapes_.size()) |
There was a problem hiding this comment.
这个检查可以提前到 ISend::Forward() 的入口处，在forward 之前就先检查
There was a problem hiding this comment.
把这个检查提到了 ISend::Forward() 入口处
There was a problem hiding this comment.
这里潜在的默认行为是 IRecv::Backward send 的类型由 grad_outputs 中各个 Tensor 自身的 dtype 决定。但比如LinearBackwardInput 的 grad_outputs 可能为 FP32，但用 BF16 计算，就会变成
前向 activation：BF16
↓
下游 Linear backward 产生梯度：FP32
↓
IRecv::Backward 原样发送：FP32
↓
上游 ISend::Backward 按保存的前向类型接收：BF16
如果要解决这个问题就得在 IRecv::Backward 和 ISend::Backward 配套，都先发送梯度的类型，再发送梯度。但这样其实增加了通信成本。可以先这样过渡，留个 fixme 后续对照 Megatron 支持配置 pipeline_dtype 或类似的功能。
7e159f2 to
8f81875
Compare
fa16795 to
6cf4962
Compare
| auto precision_config = | ||
| utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); | ||
| nn::parallel::ModelParallelConfig model_parallel_config; | ||
| model_parallel_config.pipeline_dtype = ParseDtype(FLAGS_dtype); |
There was a problem hiding this comment.
mixtral 的 main.cc 是不是也得类似改一下，或者加一个 TODO？
There was a problem hiding this comment.
mixtral 也进行了修改
| const std::unordered_set<std::string> kSupportedLRDecayStyles = { | ||
| "none", "constant", "linear", "cosine", "inverse-square-root"}; | ||
|
|
||
| DataType ParseDtype(const std::string &dtype) { |
There was a problem hiding this comment.
这部分以及kDtypeFP32[] 能不能抽到 datatype.h/cc 里
| // Dtype shared by activation and gradient payloads on the pipeline wire. | ||
| DataType pipeline_dtype_ = DataType::kFLOAT32; | ||
| // Forward tensors converted to the pipeline dtype before sending. | ||
| std::vector<std::shared_ptr<Tensor>> wire_tensors_; |
There was a problem hiding this comment.
wire_tensors_ 和 wire_grad_tensors_ 的生命周期现在一直到反向结束，能不能在使用结束之后释放掉
There was a problem hiding this comment.
不再ISend IRecv 一直持有， 改为局部变量，只用于Forward Backward的发送接收
6cf4962 to
6acd897
Compare
6acd897 to
be3e354
Compare
背景
Pipeline Parallel 在 stage 间接收 activation 和 gradient 时，固定使用
float32创建接收张量。当训练启用bfloat16等非 FP32 dtype 时，通信两端的 buffer dtype 可能不一致，导致 PP 通信失败或程序崩溃。修改内容
ModelParallelConfig中引入pipeline_dtype，作为 PP stage 间统一的通信 dtype，语义对齐 Megatron。--dtype映射为pipeline_dtype。pipeline_dtype；接收端使用pipeline_dtype创建 buffer，接收后直接交给当前 stage，由既有 autocast 机制处理算子级 dtype。pipeline_dtype；上游使用pipeline_dtype创建接收 buffer，并在接收后恢复为前向原始 input dtype。