bridge.recipes.nemotronh.nemotronh#

Module Contents#

Classes#

NemotronHCommonKwargs

Typed options accepted by NemotronH recipe helper functions.

Functions#

nemotronh_4b_pretrain_config

Return a pre-training config for NemotronH 4B.

nemotronh_8b_pretrain_config

Return a pre-training config for NemotronH 8B.

nemotronh_47b_pretrain_config

Return a pre-training config for NemotronH 47B.

nemotronh_56b_pretrain_config

Return a pre-training config for NemotronH 56B.

_nemotronh_common

Create a pre-training configuration for NemotronH models.

API#

class bridge.recipes.nemotronh.nemotronh.NemotronHCommonKwargs#

Bases: typing_extensions.TypedDict

Typed options accepted by NemotronH recipe helper functions.

Initialization

Initialize self. See help(type(self)) for accurate signature.

model_provider: megatron.bridge.models.nemotronh.NemotronHModel4BProvider | megatron.bridge.models.nemotronh.NemotronHModel8BProvider | megatron.bridge.models.nemotronh.NemotronHModel47BProvider | megatron.bridge.models.nemotronh.NemotronHModel56BProvider#

None

tokenizer_model: str | None#

None

dir: str | None#

None

name: str#

None

data_paths: list[str] | None#

None

data_args_path: str | None#

None

train_data_path: list[str] | None#

None

valid_data_path: list[str] | None#

None

test_data_path: list[str] | None#

None

per_split_data_args_path: str | None#

None

mock: bool#

None

tensor_parallelism: int#

None

pipeline_parallelism: int#

None

pipeline_parallelism_dtype: torch.dtype | None#

None

virtual_pipeline_parallelism: int | None#

None

context_parallelism: int#

None

sequence_parallelism: bool#

None

train_iters: int#

None

global_batch_size: int#

None

micro_batch_size: int#

None

seq_length: int#

None

lr: float#

None

min_lr: float#

None

lr_warmup_iters: int#

None

lr_decay_iters: int | None#

None

use_null_tokenizer: bool#

None

precision_config: megatron.bridge.training.mixed_precision.MixedPrecisionConfig | str | None#

None

comm_overlap_config: megatron.bridge.training.comm_overlap.CommOverlapConfig | None#

None

enable_default_comm_overlap: bool#

None

bridge.recipes.nemotronh.nemotronh.nemotronh_4b_pretrain_config(
**user_kwargs: typing_extensions.Unpack[bridge.recipes.nemotronh.nemotronh.NemotronHCommonKwargs],
) → megatron.bridge.training.config.ConfigContainer#

Return a pre-training config for NemotronH 4B.

This recipe is designed for single-node training (1 node). Default parallelism: TP=1, PP=1, SP=False.

See _nemotronh_common for the full list of parameters.

bridge.recipes.nemotronh.nemotronh.nemotronh_8b_pretrain_config(
**user_kwargs: typing_extensions.Unpack[bridge.recipes.nemotronh.nemotronh.NemotronHCommonKwargs],
) → megatron.bridge.training.config.ConfigContainer#

Return a pre-training config for NemotronH 8B.

This recipe is designed for single-node training (1 node). Default parallelism: TP=2, PP=1, SP=True.

See _nemotronh_common for the full list of parameters.

bridge.recipes.nemotronh.nemotronh.nemotronh_47b_pretrain_config(
**user_kwargs: typing_extensions.Unpack[bridge.recipes.nemotronh.nemotronh.NemotronHCommonKwargs],
) → megatron.bridge.training.config.ConfigContainer#

Return a pre-training config for NemotronH 47B.

This recipe is designed for single-node training (1 node with 8 GPUs). Default parallelism: TP=8, PP=1, SP=True.

Note: Uses FP8 precision by default. Communication overlap is disabled by default due to known issues with FP8 current scaling.

See _nemotronh_common for the full list of parameters.

bridge.recipes.nemotronh.nemotronh.nemotronh_56b_pretrain_config(
**user_kwargs: typing_extensions.Unpack[bridge.recipes.nemotronh.nemotronh.NemotronHCommonKwargs],
) → megatron.bridge.training.config.ConfigContainer#

Return a pre-training config for NemotronH 56B.

This recipe is designed for single-node training (1 node with 8 GPUs). Default parallelism: TP=8, PP=1, SP=True.

Note: Uses FP8 precision by default. Communication overlap is disabled by default due to known issues with FP8 current scaling.

See _nemotronh_common for the full list of parameters.

bridge.recipes.nemotronh.nemotronh._nemotronh_common(
model_provider: type[megatron.bridge.models.nemotronh.NemotronHModel4BProvider] | type[megatron.bridge.models.nemotronh.NemotronHModel8BProvider] | type[megatron.bridge.models.nemotronh.NemotronHModel47BProvider] | type[megatron.bridge.models.nemotronh.NemotronHModel56BProvider],
tokenizer_model: str | None = None,
dir: str | None = None,
name: str = 'default',
data_paths: list[str] | None = None,
data_args_path: str | None = None,
train_data_path: list[str] | None = None,
valid_data_path: list[str] | None = None,
test_data_path: list[str] | None = None,
per_split_data_args_path: str | None = None,
mock: bool = False,
tensor_parallelism: int = 1,
pipeline_parallelism: int = 1,
pipeline_parallelism_dtype: torch.dtype | None = torch.bfloat16,
virtual_pipeline_parallelism: int | None = None,
context_parallelism: int = 1,
sequence_parallelism: bool = False,
train_iters: int = 1168251,
global_batch_size: int = 768,
micro_batch_size: int = 1,
seq_length: int = 8192,
lr: float = 0.0003,
min_lr: float = 3e-05,
lr_warmup_iters: int = 2000,
lr_decay_iters: int | None = None,
use_null_tokenizer: bool = True,
precision_config: megatron.bridge.training.mixed_precision.MixedPrecisionConfig | str | None = 'bf16_mixed',
comm_overlap_config: megatron.bridge.training.comm_overlap.CommOverlapConfig | None = None,
enable_default_comm_overlap: bool = True,
) → megatron.bridge.training.config.ConfigContainer#

Create a pre-training configuration for NemotronH models.

Parameters:

  • model_provider – The model provider class for the specific NemotronH variant.

  • tokenizer_model – HuggingFace tokenizer model name (only used when use_null_tokenizer=False).

  • dir – Base directory for saving logs and checkpoints.

  • name – Name of the pre-training run.

  • data_paths – List of paths to dataset files. If None, mock data will be used.

  • data_args_path – Path to file containing data arguments.

  • train_data_path – List of training data paths.

  • valid_data_path – List of validation data paths.

  • test_data_path – List of test data paths.

  • per_split_data_args_path – Path to JSON file with per-split data configuration.

  • mock – Whether to use mock data. If True, ignores data_paths.

  • tensor_parallelism – Degree of tensor model parallelism.

  • pipeline_parallelism – Degree of pipeline model parallelism.

  • pipeline_parallelism_dtype – Data type for pipeline parallelism.

  • virtual_pipeline_parallelism – Size of virtual pipeline parallelism.

  • context_parallelism – Degree of context parallelism to be passed to model_config.

  • sequence_parallelism – Whether to use sequence parallelism.

  • train_iters – Total number of training iterations.

  • global_batch_size – Global batch size for training.

  • micro_batch_size – Micro batch size for training.

  • seq_length – Sequence length for training data.

  • lr – Learning rate.

  • min_lr – Minimum learning rate for cosine decay.

  • lr_warmup_iters – Number of warmup iterations for the learning rate.

  • lr_decay_iters – Number of iterations for learning rate decay.

  • use_null_tokenizer – Whether to use NullTokenizer instead of HuggingFaceTokenizer.

  • precision_config – Precision configuration for the model.

  • comm_overlap_config – Communication overlap configuration for the model.

  • enable_default_comm_overlap – Whether to enable default comm overlap config if none is provided.

Returns:

Configuration for pre-training.

Return type:

ConfigContainer