Make NVTE tensor handle pool size configurable

[lhb8125]

Summary

• Add runtime environment variables to configure the internal NVTETensor and NVTEGroupedTensor handle pool sizes.
• Keep the default pool size at 20 MiB to preserve existing behavior.
• Improve exhaustion errors to mention the configured pool size and the env var to increase.

Motivation

Large model initialization paths can legitimately create more TE tensor handles than the current fixed-size pool allows, even when GPU and CPU memory are otherwise sufficient. Exposing the pool size as an environment variable avoids downstream source patches for these scale-dependent cases.

Testing

• git diff --check

[greptile-apps]

Greptile Summary

This PR makes the internal NVTETensor and NVTEGroupedTensor handle pool sizes configurable via the environment variables NVTE_TENSOR_HANDLE_POOL_SIZE_MB and NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB, while preserving the existing 20 MiB default. The error messages on pool exhaustion are also improved to include the current pool size and the relevant env var.

• Adds GetTensorHandlePoolSizeMB with full input validation: whitespace trimming, digit-only enforcement, zero rejection, and errno/overflow guards — directly addressing the partial-parse and sign-wrapping problems from earlier review rounds.
• Adds GetTensorHandlePoolCapacity to compute the handle count from the validated size, and wires both helpers into TensorAllocator and GroupedTensorAllocator as const member initializers in correct declaration order.
• Documents both env vars in docs/envvars.rst under a new "General" subsection.

Confidence Score: 5/5

The change is safe to merge — it is backward-compatible (default preserved at 20 MiB), the new validation logic is thorough, and the member initializer ordering in both allocator classes is correct.

The input validation covers the full set of bad-input cases (whitespace-only, non-digit prefix, partial-parse trailing chars, zero, ERANGE overflow, SIZE_MAX conversion overflow), all previously flagged issues are resolved, and the singleton construction order is sound. No logic errors were found.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/transformer_engine.cpp	Adds configurable pool-size helpers with solid input validation (whitespace trim, digit-only guard, ERANGE + overflow checks); member initializer ordering in both allocator classes is correct; no logic errors found.
docs/envvars.rst	New env-var entries placed under a new General subsection; type, default, and description are accurate and consistent with the implementation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["TensorAllocator / GroupedTensorAllocator\n(singleton construction)"] --> B["GetTensorHandlePoolSizeMB(env_var)"]
    B --> C{env var set\nand non-empty?}
    C -- No --> D["Return default\n(20 MiB)"]
    C -- Yes --> E[Trim whitespace]
    E --> F{First char\na digit?}
    F -- No --> G["NVTE_CHECK fail:\nmust be positive integer"]
    F -- Yes --> H["strtoul(begin, &end, 10)"]
    H --> I{end == expected_end\nAND value > 0?}
    I -- No --> G
    I -- Yes --> J{errno == ERANGE?}
    J -- Yes --> K["NVTE_CHECK fail:\ntoo large"]
    J -- No --> L{pool_size_mb ≤\nSIZE_MAX / kBytesPerMB?}
    L -- No --> K
    L -- Yes --> M["Return validated\npool size (MiB)"]
    D --> N["GetTensorHandlePoolCapacity\n(pool_size_mb, sizeof(Handle))"]
    M --> N
    N --> O{pool_size_bytes ≥\nhandle_size?}
    O -- No --> P["NVTE_CHECK fail:\npool too small for 1 handle"]
    O -- Yes --> Q["Return pool_size_bytes / handle_size\n(MAX_TENSOR_NUM)"]
    Q --> R["memory.reserve(MAX_TENSOR_NUM)"]

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"}}}%%
flowchart TD
    A["TensorAllocator / GroupedTensorAllocator\n(singleton construction)"] --> B["GetTensorHandlePoolSizeMB(env_var)"]
    B --> C{env var set\nand non-empty?}
    C -- No --> D["Return default\n(20 MiB)"]
    C -- Yes --> E[Trim whitespace]
    E --> F{First char\na digit?}
    F -- No --> G["NVTE_CHECK fail:\nmust be positive integer"]
    F -- Yes --> H["strtoul(begin, &end, 10)"]
    H --> I{end == expected_end\nAND value > 0?}
    I -- No --> G
    I -- Yes --> J{errno == ERANGE?}
    J -- Yes --> K["NVTE_CHECK fail:\ntoo large"]
    J -- No --> L{pool_size_mb ≤\nSIZE_MAX / kBytesPerMB?}
    L -- No --> K
    L -- Yes --> M["Return validated\npool size (MiB)"]
    D --> N["GetTensorHandlePoolCapacity\n(pool_size_mb, sizeof(Handle))"]
    M --> N
    N --> O{pool_size_bytes ≥\nhandle_size?}
    O -- No --> P["NVTE_CHECK fail:\npool too small for 1 handle"]
    O -- Yes --> Q["Return pool_size_bytes / handle_size\n(MAX_TENSOR_NUM)"]
    Q --> R["memory.reserve(MAX_TENSOR_NUM)"]

Loading

_{Reviews (4): Last reviewed commit: "Simplify tensor handle pool env parsing" | Re-trigger Greptile}

[ptrendx]

I am not opposed to creating such a variable, but I would really like to see an example of such legitimate use which goes over this limit. Could you run the experiment that is failing for you with https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/common/transformer_engine.cpp#L487 set to true and send me the log of that?

[lhb8125]

I am not opposed to creating such a variable, but I would really like to see an example of such legitimate use which goes over this limit. Could you run the experiment that is failing for you with https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/common/transformer_engine.cpp#L487 set to true and send me the log of that?

te_handle_pool_debug_fakepg_20260616-010629_2134338.pruned.log
I attached a pruned version, which replaces the repeated lines with a comment. The case is EP1 for Qwen3-235B training, so the number of local experts is 128, where we adopt full recompute to control the dynamic activation's memory footprint.
I sent the whole log to you via slack.

[ptrendx]

Why not just use std::atoi?

[lhb8125]

I avoided std::atoi here because it does not give us enough validation/error handling for an env var that controls an allocation size. In particular, it cannot distinguish parse failure from 0, accepts prefixes such as "50bad", returns int rather than size_t, and gives poor behavior/diagnostics for negative or overflowing inputs.

The current parser is a bit more verbose, but it intentionally enforces a full-string positive integer and checks overflow before converting to bytes, so invalid values fail with a clear NVTE_CHECK message.

If you prefer using a standard parser, I can switch this to std::strtoull/std::from_chars with explicit end-pointer/overflow checks, but I would avoid std::atoi for this case.

[ptrendx]

Considering that it is supposed to take the value in MBs, at least some of the concerns here are not really relevant - e.g. I do not expect anybody to have a problem with int vs size_t here as the values would be < 1000. I'm fine with strtoul (again, we don't need ull here) - basically trying to make sure that the code is as understandable as possible.

[lhb8125]

Thanks, agreed. I simplified this to use std::strtoul and kept the validation intentionally small:

• trim surrounding whitespace
• require the first non-space character to be a digit
• require strtoul to consume the full trimmed value
• require the parsed value to be > 0
• check ERANGE for overflow

[ptrendx]

Ok, I see why it happens. I'm ok with this PR (modulo one comment about string to integer conversion). Now, the real reason why this happens, while legitimate, is also very bad performance-wise (logic for the master weight casting to lower precision) and this should also be fixed in a separate PR.