From 5c47ef51307233ef76bcfe53775bd2d024230f3c Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 24 Mar 2026 21:30:47 -0700 Subject: [PATCH] adam: separate param/grad dtypes in AdamCapturableFunctor Signed-off-by: Ritesh Patel --- .../common/multi_tensor/adam.cu | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/transformer_engine/common/multi_tensor/adam.cu b/transformer_engine/common/multi_tensor/adam.cu index b2c8f29a13..331d6230a7 100644 --- a/transformer_engine/common/multi_tensor/adam.cu +++ b/transformer_engine/common/multi_tensor/adam.cu @@ -380,7 +380,7 @@ struct AdamFunctor { } }; -template +template struct AdamCapturableFunctor { __device__ __forceinline__ void operator()(int64_t chunk_size, volatile int *noop_gmem, TensorListMetadata<4> &tl, // NOLINT(*) @@ -404,10 +404,10 @@ struct AdamCapturableFunctor { int chunk_idx = tl.block_to_chunk[blockIdx.x]; int64_t n = tl.sizes[tensor_loc]; - T *g = reinterpret_cast(tl.addresses[0][tensor_loc]); + GRAD_T *g = reinterpret_cast(tl.addresses[0][tensor_loc]); g += chunk_idx * chunk_size; - T *p = reinterpret_cast(tl.addresses[1][tensor_loc]); + PARAM_T *p = reinterpret_cast(tl.addresses[1][tensor_loc]); p += chunk_idx * chunk_size; MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); @@ -429,7 +429,7 @@ struct AdamCapturableFunctor { int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) { r_g[ii] = static_cast(g[i]) * (*inv_scale); - g[i] = static_cast(r_g[ii]); + g[i] = static_cast(r_g[ii]); r_p[ii] = static_cast(p[i]); r_m[ii] = static_cast(m[i]); r_v[ii] = static_cast(v[i]); @@ -465,7 +465,7 @@ struct AdamCapturableFunctor { for (int ii = 0; ii < ILP; ii++) { int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) { - p[i] = static_cast(r_p[ii]); + p[i] = static_cast(r_p[ii]); m[i] = static_cast(r_m[ii]); v[i] = static_cast(r_v[ii]); } @@ -874,13 +874,14 @@ void multi_tensor_adam_capturable_cuda(int chunk_size, Tensor noop_flag, // Check tensor dtypes const auto g_in_type_te = tensor_lists[0][0]->dtype(); + const auto p_in_type_te = tensor_lists[1][0]->dtype(); for (size_t j = 0; j < num_tensors_per_list; j++) { NVTE_CHECK(tensor_lists[0][j]->dtype() == g_in_type_te, "Grad tensor ", j, " has dtype=", to_string(tensor_lists[0][j]->dtype()), ", but expected dtype=", to_string(g_in_type_te)); - NVTE_CHECK(tensor_lists[1][j]->dtype() == g_in_type_te, "Param tensor ", j, + NVTE_CHECK(tensor_lists[1][j]->dtype() == p_in_type_te, "Param tensor ", j, " has dtype=", to_string(tensor_lists[1][j]->dtype()), - ", but expected dtype=", to_string(g_in_type_te)); + ", but expected dtype=", to_string(p_in_type_te)); { const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; @@ -899,14 +900,17 @@ void multi_tensor_adam_capturable_cuda(int chunk_size, Tensor noop_flag, // Launch kernel TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - tensor_lists[0][0]->dtype(), dtype, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( - moment_type_te, moment_type, - multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, - AdamCapturableFunctor(), stream, beta1, - beta2, reinterpret_cast(step.data.dptr), bias_correction, - epsilon, reinterpret_cast(lr.data.dptr), (adamMode_t)mode, - weight_decay, reinterpret_cast(inv_scale.data.dptr));)) + p_in_type_te, p_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + g_in_type_te, g_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<4>( + BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + AdamCapturableFunctor(), stream, beta1, + beta2, reinterpret_cast(step.data.dptr), bias_correction, epsilon, + reinterpret_cast(lr.data.dptr), (adamMode_t)mode, weight_decay, + reinterpret_cast(inv_scale.data.dptr));))) NVTE_CHECK_CUDA(cudaGetLastError()); }