bridge.perf_recipes._common

Shared helpers for flat performance benchmark recipes.

_benchmark_common applies throughput-measurement defaults. _perf_precision returns a mixed-precision config for a given dtype.

Module Contents

Functions

_benchmark_common	Apply benchmark-mode defaults that prioritize throughput measurement over convergence.
_enable_ncclep	Switch a flex-dispatcher recipe to the NCCL EP backend (GB300 MoE default).
_perf_precision	Return mixed-precision config tuned for perf benchmarks.

API

bridge.perf_recipes._common._benchmark_common(

cfg: megatron.bridge.training.config.ConfigContainer,

cross_entropy_impl: str = 'te',

) → None

Apply benchmark-mode defaults that prioritize throughput measurement over convergence.

Intended for performance benchmark recipes only. Sets short training runs, disables checkpointing/eval, tunes scheduler, and enables perf-oriented kernels.

This is the fixed-model-shape policy for flat performance recipes. Canonical recipes launched with scripts/performance --use_recipes retain their own tokenizer vocabulary policy.

Individual recipes may override any of these after calling this function (e.g. Kimi K2 sets grad_reduce_in_fp32 = True).

bridge.perf_recipes._common._enable_ncclep(

cfg: megatron.bridge.training.config.ConfigContainer,

) → None

Switch a flex-dispatcher recipe to the NCCL EP backend (GB300 MoE default).

Call after _benchmark_common. Only the dispatch stack changes; parallelism, precision, recompute and the CUDA graph mode stay as the recipe set them. The calling recipe still declares its own cfg.env_vars inline and adds "NCCL_EP_HT_EM_PULL_PUSH": 1 there.

bridge.perf_recipes._common._perf_precision(compute_dtype: str)

Return mixed-precision config tuned for perf benchmarks.

Identical to scripts/performance/utils/precision.get_precision_config but importable from the recipe package. Always sets grad_reduce_in_fp32=False so that callers that replace cfg.mixed_precision after _benchmark_common() still get the benchmark-mode default.