已合并
[fix]:solve the nan error of recompute_w_u_fwd. #3360
LinShua创建于 3月30日
[fix]:solve the nan error of recompute_w_u_fwd. #3360
已合并
共 1 个文件变更+4-3
| @@ -158,7 +158,7 @@ def recompute_w_u_fwd_kernel( | |||
| 158 | gk, | 158 | gk, |
| 159 | cu_seqlens, | 159 | cu_seqlens, |
| 160 | chunk_indices, | 160 | chunk_indices, |
| 161 | - T, | 161 | + T_tmp, |
L | |||
| 162 | B, | 162 | B, |
| 163 | H: tl.constexpr, | 163 | H: tl.constexpr, |
| 164 | K: tl.constexpr, | 164 | K: tl.constexpr, |
| @@ -173,7 +173,7 @@ def recompute_w_u_fwd_kernel( | |||
| 173 | ): | 173 | ): |
| 174 | core_id = tl.program_id(0) | 174 | core_id = tl.program_id(0) |
| 175 | total_cores = tl.num_programs(0) | 175 | total_cores = tl.num_programs(0) |
| 176 | - T_max = T | 176 | + T_max = T_tmp |
| 177 | 177 | ||
| 178 | base_chunks_per_pid = NT // total_cores | 178 | base_chunks_per_pid = NT // total_cores |
| 179 | remainder_chunks = NT % total_cores | 179 | remainder_chunks = NT % total_cores |
| @@ -195,6 +195,7 @@ def recompute_w_u_fwd_kernel( | |||
| 195 | offset = bos + i_h * T_max | 195 | offset = bos + i_h * T_max |
| 196 | T = eos - bos | 196 | T = eos - bos |
| 197 | else: | 197 | else: |
| 198 | + T = T_tmp | ||
| 198 | i_t = idx | 199 | i_t = idx |
| 199 | bos, eos = i_b * T, i_b * T + T | 200 | bos, eos = i_b * T, i_b * T + T |
| 200 | offset = bos * H + i_h * T_max | 201 | offset = bos * H + i_h * T_max |
| @@ -264,7 +265,7 @@ def recompute_w_u_fwd( | |||
| 264 | gk=gk, | 265 | gk=gk, |
| 265 | cu_seqlens=cu_seqlens, | 266 | cu_seqlens=cu_seqlens, |
| 266 | chunk_indices=chunk_indices, | 267 | chunk_indices=chunk_indices, |
| 267 | - T=T, | 268 | + T_tmp=T, |
| 268 | B=B, | 269 | B=B, |
| 269 | H=H, | 270 | H=H, |
| 270 | K=K, | 271 | K=K, |
这种tmp， 不太建议