core.transformer.cuda_graph_config

Module Contents

Functions

normalize_cuda_graph_modules	Normalize mixed CUDA graph scope inputs into enum values plus deprecation metadata.
normalize_inference_cuda_graph_scope	Normalize inference CUDA graph scope and apply the impl-derived default.
validate_deprecated_cuda_graph_modules_migration_inputs	Reject ambiguous mixed old/new CUDA graph inputs before applying migration.
get_deprecated_cuda_graph_modules_migration	Return the effective new-style migration for a deprecated cuda_graph_modules value.

Data

CUDA_GRAPH_MODULES_DEPRECATIONS
ALLOWED_INFERENCE_SCOPES

API

core.transformer.cuda_graph_config.CUDA_GRAPH_MODULES_DEPRECATIONS

None

core.transformer.cuda_graph_config.ALLOWED_INFERENCE_SCOPES: dict[str, Set[megatron.core.transformer.enums.InferenceCudaGraphScope]]

None

core.transformer.cuda_graph_config.normalize_cuda_graph_modules(

scopes: Optional[Union[str, megatron.core.transformer.enums.CudaGraphModule, List[Union[str, megatron.core.transformer.enums.CudaGraphModule]]]],

) → Tuple[List[megatron.core.transformer.enums.CudaGraphModule], List[Tuple[str, str, object]], bool]

Normalize mixed CUDA graph scope inputs into enum values plus deprecation metadata.

core.transformer.cuda_graph_config.normalize_inference_cuda_graph_scope(

scope: Optional[Union[str, megatron.core.transformer.enums.InferenceCudaGraphScope]],

cuda_graph_impl: str,

) → megatron.core.transformer.enums.InferenceCudaGraphScope

Normalize inference CUDA graph scope and apply the impl-derived default.

core.transformer.cuda_graph_config.validate_deprecated_cuda_graph_modules_migration_inputs(

deprecated_scopes: List[Tuple[str, str, object]],

cuda_graph_impl: str,

inference_cuda_graph_scope: Optional[Union[str, megatron.core.transformer.enums.InferenceCudaGraphScope]],

) → None

Reject ambiguous mixed old/new CUDA graph inputs before applying migration.

Deprecated scope strings are still accepted for compatibility, but only when they are not combined with conflicting new-style fields.

core.transformer.cuda_graph_config.get_deprecated_cuda_graph_modules_migration(

scope: str,

attr: str,

value: object,

cuda_graph_impl: str,

) → Optional[Tuple[str, object]]

Return the effective new-style migration for a deprecated cuda_graph_modules value.