adam: separate param/grad dtypes in AdamCapturableFunctor

[rapatel]

Description

Split AdamCapturableFunctor's param/grad dtypes and relax host-side validation, to match the non-capturable Adam behavior.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Split AdamCapturableFunctor's param/grad dtypes and relax host-side validation, to match the non-capturable Adam behavior.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR fixes a restriction in AdamCapturableFunctor where grad and param tensors were required to share the same dtype, aligning the capturable path's behavior with the existing non-capturable AdamFunctor which already supports separate PARAM_T / GRAD_T template parameters.

• Kernel template: splits the single T type parameter in AdamCapturableFunctor into PARAM_T and GRAD_T, updating all pointer casts and write-backs accordingly.
• Host validation: adds p_in_type_te derived from tensor_lists[1][0], checks each param tensor against it independently (previously checked against g_in_type_te), and nests a second TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY to dispatch across all param×grad×moment type combinations (3×3×2 = 18 instantiations vs. the prior 6).

Confidence Score: 4/5

The change is narrow and internally consistent: it mirrors an existing, tested design in the non-capturable AdamFunctor. The only functional gap is that AdamCapturableMasterFunctor retains the old coupled-type constraint, but that path has a different architecture and the PR does not claim to fix it.

The kernel change is straightforward and the host validation is updated correctly. The nested type switches produce 3×3×2 = 18 kernel instantiations instead of 6, which is expected overhead for supporting independent param/grad dtypes. No tests are added, and AdamCapturableMasterFunctor retains the coupled dtype assumption — both are minor gaps worth tracking but do not affect the correctness of the capturable non-master path being fixed here.

transformer_engine/common/multi_tensor/adam.cu — specifically AdamCapturableMasterFunctor and multi_tensor_adam_capturable_master_cuda, which still enforce identical grad/param dtype and were not updated in this PR.

Important Files Changed

Filename	Overview
transformer_engine/common/multi_tensor/adam.cu	Splits AdamCapturableFunctor's single T into PARAM_T/GRAD_T, fixes host-side dtype validation, and adds nested type dispatch; logic is correct and mirrors the non-capturable AdamFunctor. No tests added and AdamCapturableMasterFunctor is not updated with the same split.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant multi_tensor_adam_capturable_cuda
    participant TypeSwitch as Type Switch Macros
    participant Kernel as AdamCapturableFunctor

    Caller->>multi_tensor_adam_capturable_cuda: tensor_lists, lr, beta1, beta2, step, inv_scale
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: Validate 4 tensor lists (g, p, m, v)
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "g_in_type_te = lists[0][0].dtype()"
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "p_in_type_te = lists[1][0].dtype()"
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "Check all grads == g_in_type_te"
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "Check all params == p_in_type_te"
    multi_tensor_adam_capturable_cuda->>TypeSwitch: "SWITCH(p_in_type_te -> PARAM_T)"
    TypeSwitch->>TypeSwitch: "SWITCH(g_in_type_te -> GRAD_T)"
    TypeSwitch->>TypeSwitch: "SWITCH(moment_type -> MOMENT_T)"
    TypeSwitch->>Kernel: multi_tensor_apply with AdamCapturableFunctor
    Kernel->>Kernel: "r_g = cast float(g[i]) * inv_scale, write back g[i]"
    Kernel->>Kernel: "r_p = cast float(p[i])"
    Kernel->>Kernel: Adam math in float
    Kernel->>Kernel: "p[i] = cast PARAM_T(r_p), m/v written back"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant multi_tensor_adam_capturable_cuda
    participant TypeSwitch as Type Switch Macros
    participant Kernel as AdamCapturableFunctor

    Caller->>multi_tensor_adam_capturable_cuda: tensor_lists, lr, beta1, beta2, step, inv_scale
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: Validate 4 tensor lists (g, p, m, v)
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "g_in_type_te = lists[0][0].dtype()"
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "p_in_type_te = lists[1][0].dtype()"
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "Check all grads == g_in_type_te"
    multi_tensor_adam_capturable_cuda->>multi_tensor_adam_capturable_cuda: "Check all params == p_in_type_te"
    multi_tensor_adam_capturable_cuda->>TypeSwitch: "SWITCH(p_in_type_te -> PARAM_T)"
    TypeSwitch->>TypeSwitch: "SWITCH(g_in_type_te -> GRAD_T)"
    TypeSwitch->>TypeSwitch: "SWITCH(moment_type -> MOMENT_T)"
    TypeSwitch->>Kernel: multi_tensor_apply with AdamCapturableFunctor
    Kernel->>Kernel: "r_g = cast float(g[i]) * inv_scale, write back g[i]"
    Kernel->>Kernel: "r_p = cast float(p[i])"
    Kernel->>Kernel: Adam math in float
    Kernel->>Kernel: "p[i] = cast PARAM_T(r_p), m/v written back"

Loading

Comments Outside Diff (1)

1. transformer_engine/common/multi_tensor/adam.cu, line 940-942 (link)

   AdamCapturableMasterFunctor param/grad dtype still coupled

   AdamCapturableMasterFunctor (line 477) still uses a single T for both g and p, and multi_tensor_adam_capturable_master_cuda (line 940) still validates param dtype against g_in_type_te. If the same mixed-dtype use-case is needed for the master-weight capturable path, it would need an equivalent split there. This is likely out of scope for this PR, but worth documenting so it isn't overlooked later.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "adam: separate param/grad dtypes in Adam..." | Re-trigger Greptile}

[vthumbe1503]

LGTM @rapatel. Could we also add a test in test_fused_optimizer.py for the case where param and grad have different types and verify the results?