maxtext.configs.pyconfig_deprecated module

Contents

maxtext.configs.pyconfig_deprecated module#

maxtext.configs.pyconfig_deprecated.yaml_key_to_env_key(s)[source]#
Parameters:

s (str)

Return type:

str

maxtext.configs.pyconfig_deprecated.string_to_bool(s)[source]#
Parameters:

s (str)

Return type:

bool

maxtext.configs.pyconfig_deprecated.validate_compute_axis_order(s)[source]#
Parameters:

s (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_shard_mode(shard_mode, decoder_block, quantization)[source]#

Validates sharding settings, raising ValueError for incompatible combinations.

Parameters:
  • shard_mode (str)

  • decoder_block (str)

  • quantization (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_kv_quant_axis(s, quantize_kvcache)[source]#
Parameters:
  • s (str)

  • quantize_kvcache (bool)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_attention_kernel(s)[source]#
Parameters:

s (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_attention_type(s)[source]#
Parameters:

s (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_moba_attention(moba, attention)[source]#
Return type:

None

maxtext.configs.pyconfig_deprecated.validate_attention_window_params(attention_type, chunk_attn_window_size, sliding_window_size)[source]#

Validates window size parameters for attention types ‘chunk’ and ‘local’.

Parameters:
  • attention_type (str)

  • chunk_attn_window_size (int)

  • sliding_window_size (int)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_profiler_type(s)[source]#
Parameters:

s (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_periodic_profiler(profiler, profile_periodically_period, profiler_steps)[source]#
maxtext.configs.pyconfig_deprecated.validate_model_call_mode(s)[source]#
Parameters:

s (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_prefill_and_target_lengths(max_prefill_length, max_target_length)[source]#
Parameters:
  • max_prefill_length (int)

  • max_target_length (int)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_rope_type(rope_type)[source]#
Parameters:

rope_type (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_expert_shard_attention_option(expert_shard_attention_option)[source]#
Parameters:

expert_shard_attention_option (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_vocab_tiling(num_vocab_tiling, per_device_batch_size, max_target_length, enable_nnx)[source]#
Parameters:
  • num_vocab_tiling (int)

  • per_device_batch_size (int)

  • max_target_length (int)

  • enable_nnx (bool)

maxtext.configs.pyconfig_deprecated.validate_rampup_batch_size(batch_size_start, batch_size_end, batch_size_increment, global_rampup_samples)[source]#
maxtext.configs.pyconfig_deprecated.validate_context_parallel_strategy_ring(context_parallel_size, context_parallel_strategy, hardware)[source]#

Validates that ring context parallelism strategy is only used on GPU hardware.

Parameters:
  • context_parallel_size (int)

  • context_parallel_strategy (str)

  • hardware (str)

Return type:

None

maxtext.configs.pyconfig_deprecated.validate_keys(keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_tokenizer(keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_constant_bound(keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_quantization_methods(keys)[source]#

Validate quantization methods

maxtext.configs.pyconfig_deprecated.validate_tokamax_usage(keys)[source]#

Validate tokamax usage for gmm kernel

maxtext.configs.pyconfig_deprecated.validate_data_input(keys)[source]#

validate provided parameters for data input

maxtext.configs.pyconfig_deprecated.validate_llama4_config(keys)[source]#

Validates the following checks for Llama4 models:

Parameters:

keys (dict) – the raw config in dict form

maxtext.configs.pyconfig_deprecated.validate_model_name(s)[source]#

Validate provided model name.

Parameters:

s (str)

Return type:

bool

maxtext.configs.pyconfig_deprecated.validate_multimodal_model_name(s)[source]#
Parameters:

s (str)

Return type:

bool

maxtext.configs.pyconfig_deprecated.validate_no_keys_overwritten_twice(keys1, keys2)[source]#
Parameters:
  • keys1 (list[str])

  • keys2 (list[str])

maxtext.configs.pyconfig_deprecated.validate_and_assign_remat_tensors(keys)[source]#
maxtext.configs.pyconfig_deprecated.resolve_config_path(param)[source]#

Resolve config path to auto rewrite to use new src folder. This ensures backwards compatibility with older versions of MaxText.

Parameters:

param (str)

Return type:

str

maxtext.configs.pyconfig_deprecated.create_parallelisms_list(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.set_mu_dtype(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_and_set_hlo_dump_defaults(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_multiple_slices(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.set_and_validate_pipeline_config(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_deepseek_moe(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_mlp_dim(raw_keys)[source]#

Validates that MLP dimensions are consistent for fully MoE models.

maxtext.configs.pyconfig_deprecated.validate_gpt_oss_moe(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_sparse_matmul_parallelism(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_ring_of_experts_parallelism(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_shard_expert_on_fsdp(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_ragged_dot(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.validate_optimizer_sharding_over_data(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.create_new_logical_axis_rules(old_logical_axis_rules, new_logical_axis_rules)[source]#
maxtext.configs.pyconfig_deprecated.update_model_keys(raw_keys, model_keys, key)[source]#

Update key value in raw_keys from the value in model_keys.

maxtext.configs.pyconfig_deprecated.validate_and_update_keys(raw_keys, model_keys, config_name)[source]#

Validate and update model specific config keys

Parameters:

config_name (str)

maxtext.configs.pyconfig_deprecated.get_individual_scales(scale)[source]#

Choose appropriate scales for individual dimensions based on global scale

We choose to rotate between doubling:

  • num_head and mlp_dim

  • embed_dim

  • num_layers

Any one of these steps is not a perfect doubling, although going through a cycle of three is a near perfect 8x scaling except for the linear -> softmax -> output step

maxtext.configs.pyconfig_deprecated.calculate_global_batch_sizes(per_device_batch_size, expansion_factor_real_data, num_devices, gradient_accumulation_steps)[source]#

Calculates target global batch size from target devices and per_device_batch

maxtext.configs.pyconfig_deprecated.calculate_rampup_samples_and_steps(batch_size_start, batch_size_end, batch_size_increment, global_rampup_samples)[source]#

Calculate num of samples for each increment and num of steps for batch rampup

maxtext.configs.pyconfig_deprecated.get_num_target_devices(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.get_quantization_local_shard_count(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.get_context_parallel_size(raw_keys)[source]#
maxtext.configs.pyconfig_deprecated.using_pipeline_parallelism(raw_keys)[source]#
Return type:

bool

maxtext.configs.pyconfig_deprecated.using_tensor_parallelism(raw_keys)[source]#
Return type:

bool

maxtext.configs.pyconfig_deprecated.using_sequence_parallelism(raw_keys)[source]#
Return type:

bool

maxtext.configs.pyconfig_deprecated.using_expert_parallelism(raw_keys)[source]#
Return type:

bool

maxtext.configs.pyconfig_deprecated.using_fsdp_and_transpose_parallelism(raw_keys)[source]#
Return type:

bool

class maxtext.configs.pyconfig_deprecated.HyperParameters(config)[source]#

Bases: object

Wrapper class to expose the configuration in a read-only manner.

get_keys()[source]#
tree_flatten()[source]#
classmethod tree_unflatten(aux_data, children)[source]#
maxtext.configs.pyconfig_deprecated.initialize(argv, **kwargs)[source]#