Index

_

__post_init__() (eformer.optimizers._config.SchedulerConfig method) _actor_pool (eformer.executor.ray.pool_manager.ResourcePoolManager attribute) (eformer.executor.ray.pool_manager.SlicePoolManager attribute) _count (eformer.pytree._tree_util.MetaValueRecreator attribute) _create_cosine() (eformer.optimizers._factory.SchedulerFactory method) _create_linear() (eformer.optimizers._factory.SchedulerFactory method) _last_scale_check_ts (eformer.executor.ray.pool_manager.SlicePoolManager attribute)

_last_scale_ts (eformer.executor.ray.pool_manager.SlicePoolManager attribute) _level (eformer.executor.patch_tpus_ray.LazyLogger attribute) _logger (eformer.executor.patch_tpus_ray.LazyLogger attribute) (eformer.loggings.ProgressLogger attribute) _name (eformer.executor.patch_tpus_ray.LazyLogger attribute) _rng (eformer.pytree._tree_util.MetaValueRecreator attribute) _tpu_type (eformer.executor.ray.pool_manager.SlicePoolManager attribute) _validate() (eformer.optimizers._config.SchedulerConfig method)

A

absmax (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 attribute), [1] abstract_mesh (eformer.common_types.Mesh property) accelerator_type (eformer.executor.ray.resource_manager.RayResources attribute), [1] actor (eformer.executor.ray.pool_manager.ActorPoolMember attribute), [1] actor_info (eformer.executor.ray.pool_manager.ActorPoolMember attribute), [1] ActorPoolMember (class in eformer.executor.ray.pool_manager) AdafactorConfig (class in eformer.optimizers._config) AdafactorOptimizer (class in eformer.optimizers._builders) adam_b1 (eformer.optimizers._config.MuonConfig attribute), [1] adam_b2 (eformer.optimizers._config.MuonConfig attribute), [1] adam_eps_root (eformer.optimizers._config.MuonConfig attribute), [1] adam_weight_decay (eformer.optimizers._config.MuonConfig attribute), [1] adamw_b1 (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] adamw_b2 (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] adamw_eps (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] adamw_eps_root (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] adamw_mu_dtype (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] AdamWConfig (class in eformer.optimizers._config) AdamWOptimizer (class in eformer.optimizers._builders) adaptive (eformer.optimizers._config.MuonConfig attribute), [1] add_1bit_xy() (in module eformer.ops.quantization.implicit_array_1bit) add_8bit_xy() (in module eformer.ops.quantization.implicit_array_8bit) add_nf4_xy() (in module eformer.ops.quantization.implicit_array_nf4) address (eformer.executor.cluster_util.RayClusterConfig attribute), [1] adjust() (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale method) (eformer.mpric.loss_scaling.loss_scaler.NoOpLossScale method) AMD_INSTINCT_MI100 (eformer.executor.ray.resource_manager.HardwareType attribute) AMD_INSTINCT_MI210 (eformer.executor.ray.resource_manager.HardwareType attribute) AMD_INSTINCT_MI250 (eformer.executor.ray.resource_manager.HardwareType attribute) AMD_INSTINCT_MI250x (eformer.executor.ray.resource_manager.HardwareType attribute) AMD_INSTINCT_MI300x (eformer.executor.ray.resource_manager.HardwareType attribute) AMD_RADEON_HD_7900 (eformer.executor.ray.resource_manager.HardwareType attribute) AMD_RADEON_R9_200_HD_7900 (eformer.executor.ray.resource_manager.HardwareType attribute) analyze_pytree() (eformer.escale.mesh.mesh_helpers.MeshPartitionHelper method) analyze_sharding_strategy() (in module eformer.escale.partition.constraints) apply() (eformer.escale.helpers.base.AutoShardingRule method) (eformer.escale.helpers.base.CompositeShardingRule method) (eformer.escale.helpers.base.MemoryConstrainedShardingRule method) (eformer.escale.helpers.base.ShapeBasedShardingRule method) (eformer.escale.helpers.base.ShardingRule method) apply_gather_functions() (in module eformer.serialization.utils) apply_gradients_stage_local() (eformer.optimizers._stage_local.StageLocalGradientTransformation method) apply_logical_sharding() (in module eformer.escale.partition.manager) apply_sharding_tree() (in module eformer.serialization.sharding_utils) Argu() (in module eformer.aparser._aparser) ArgumentParser (class in eformer.executor.patch_tpus_ray) Array (in module eformer.common_types) Array1B (class in eformer.ops.quantization.implicit_array_1bit) Array8B (class in eformer.ops.quantization.implicit_array_8bit) array_metadata (eformer.serialization.async_manager.CheckpointMetadata attribute), [1] ArrayNF4 (class in eformer.ops.quantization.implicit_array_nf4) as_posix() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method)

astype() (eformer.jaximus._imus.ImplicitArray method) async_remove() (in module eformer.serialization.fsspec_utils) AsyncCheckpointManager (class in eformer.serialization.async_manager) attention_dim_axis (eformer.escale.partition.manager.PartitionAxis attribute) attention_kv_dim_axis (eformer.escale.partition.manager.PartitionAxis attribute) AttnKVSharding (class in eformer.common_types) AttnQSharding (class in eformer.common_types) auto_namedsharding() (in module eformer.escale.partition.auto_spec) auto_partition_spec() (in module eformer.escale.partition.auto_spec) auto_pytree() (in module eformer.pytree._pytree) auto_ray_cluster() (in module eformer.executor.cluster_util) auto_shard_array() (in module eformer.escale.partition.auto_spec) auto_shard_pytree() (eformer.escale.mesh.mesh_helpers.MeshPartitionHelper method) auto_start_cluster (eformer.executor.cluster_util.RayClusterConfig attribute), [1] autoscale_execute() (eformer.executor.ray.executor.RayExecutor method) (eformer.executor.ray.executor.RayExecutor static method) (in module eformer.executor.ray.executor) autoscale_execute_resumable() (eformer.executor.ray.executor.RayExecutor class method) (eformer.executor.ray.executor.RayExecutor method) (in module eformer.executor.ray.executor) AutoShardingRule (class in eformer.escale.helpers.base) aux_field() (in module eformer.jaximus._imus) available_cpu_cores() (in module eformer.executor.ray.resource_manager) aval (eformer.jaximus._imus.ImplicitArray property) AWS_NEURON_CORE (eformer.executor.ray.resource_manager.HardwareType attribute) axes (eformer.common_types.AttnKVSharding attribute) (eformer.common_types.AttnQSharding attribute) (eformer.common_types.ColumnWise attribute) (eformer.common_types.DynamicShardingAxes attribute), [1] (eformer.common_types.ExpertActivations attribute) (eformer.common_types.ExpertActivationsAlt attribute) (eformer.common_types.ExpertColumnWise attribute) (eformer.common_types.ExpertColumnWiseAlt attribute) (eformer.common_types.ExpertRowWise attribute) (eformer.common_types.ExpertRowWiseAlt attribute) (eformer.common_types.ExpertTensorParallel attribute) (eformer.common_types.HiddenStateSharding attribute) (eformer.common_types.Replicated attribute) (eformer.common_types.RowWise attribute) (eformer.common_types.SColumnWise attribute) (eformer.common_types.SRowWise attribute) (eformer.common_types.UnifiedExpertColumnWise attribute) (eformer.common_types.UnifiedExpertRowWise attribute) axis (eformer.ops.quantization.implicit_array_8bit.Array8B attribute), [1] axis_names (eformer.common_types.Mesh attribute) (eformer.escale.helpers.base.AutoShardingRule attribute) (eformer.escale.helpers.base.MemoryConstrainedShardingRule attribute) axis_sizes (eformer.common_types.Mesh property) (eformer.escale.mesh.mesh_helpers.MeshPartitionHelper attribute) AxisIdxes (in module eformer.common_types) AxisNames (in module eformer.common_types) AxisType (in module eformer.common_types)

B

B (eformer.optimizers._tx.white_kron.LeafState attribute), [1] b1 (eformer.optimizers._config.AdamWConfig attribute), [1] (eformer.optimizers._config.LionConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] b2 (eformer.optimizers._config.AdamWConfig attribute), [1] (eformer.optimizers._config.LionConfig attribute), [1] backend_steps (eformer.optimizers._config.ScionConfig attribute), [1] barrier_sync() (in module eformer.escale.helpers.base) base_env (eformer.executor.ray.executor.TpuRemoteManager attribute) base_path (eformer.serialization.checkpointer.Checkpointer attribute) BATCH (in module eformer.common_types) batch_axis (eformer.escale.partition.manager.PartitionAxis attribute) beta (eformer.optimizers._config.MuonConfig attribute), [1] beta1 (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.MarsConfig attribute), [1] (eformer.optimizers._config.ScionConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] beta2 (eformer.optimizers._config.MarsConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] BIAS_HEAD_SEQ (in module eformer.common_types) bias_head_sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute) bias_key_sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute) BIAS_KV_SEQ (in module eformer.common_types) BINARY (eformer.ops.quantization._config.QuantizationType attribute) blob (eformer.paths.GCSPath property) blob_name (eformer.paths.GCSPath attribute) block_size (eformer.ops.quantization._config.QuantizationConfig attribute), [1] (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 attribute), [1] (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1]

bmm_nf4() (in module eformer.ops.quantization.quantization_functions) bmm_nf4_transpose() (in module eformer.ops.quantization.quantization_functions) broadcast_in_dim_1bit_operand() (in module eformer.ops.quantization.implicit_array_1bit) broadcast_in_dim_8bit_operand() (in module eformer.ops.quantization.implicit_array_8bit) broadcast_in_dim_nf4_operand() (in module eformer.ops.quantization.implicit_array_nf4) broadcast_tensor() (in module eformer.serialization.utils) bucket (eformer.paths.GCSPath property) bucket_name (eformer.paths.GCSPath attribute) build() (eformer.optimizers._base.OptimizerBuilder method), [1] (eformer.optimizers._base.SchedulerBuilder method), [1] (eformer.optimizers._builders.AdafactorOptimizer method) (eformer.optimizers._builders.AdamWOptimizer method) (eformer.optimizers._builders.ConstantSchedulerBuilder method) (eformer.optimizers._builders.CosineSchedulerBuilder method) (eformer.optimizers._builders.LinearSchedulerBuilder method) (eformer.optimizers._builders.LionOptimizer method) (eformer.optimizers._builders.MarsOptimizer method) (eformer.optimizers._builders.MuonOptimizer method) (eformer.optimizers._builders.QuadOptimizer method) (eformer.optimizers._builders.RMSPropOptimizer method) (eformer.optimizers._builders.SkewOptimizer method) build_and_push_docker_image() (in module eformer.executor.ray.docker_executor) build_docker_image() (in module eformer.executor.docker_tpu) build_mpmd() (eformer.optimizers._base.OptimizerBuilder method) (eformer.optimizers._builders.AdafactorOptimizer method) (eformer.optimizers._builders.AdamWOptimizer method) (eformer.optimizers._builders.LionOptimizer method) (eformer.optimizers._builders.MarsOptimizer method) (eformer.optimizers._builders.MuonOptimizer method) (eformer.optimizers._builders.QuadOptimizer method) (eformer.optimizers._builders.RMSPropOptimizer method) (eformer.optimizers._builders.SkewOptimizer method)

C

calculate_host_mesh_shape() (in module eformer.escale.mesh.creation) cancel_all_futures() (eformer.executor.ray.resource_manager.RayResources static method) cast_for_compute() (eformer.mpric.handler.precision_handler.PrecisionHandler method) cast_for_output() (eformer.mpric.handler.precision_handler.PrecisionHandler method) cast_params() (eformer.mpric.handler.precision_handler.PrecisionHandler method) checkpoint_dir (eformer.serialization.base_manager.CheckpointManager attribute) Checkpointer (class in eformer.serialization.checkpointer) CheckpointInterval (class in eformer.serialization.checkpointer) CheckpointManager (class in eformer.serialization.base_manager) CheckpointMetadata (class in eformer.serialization.async_manager) checksum (eformer.serialization.async_manager.CheckpointMetadata attribute), [1] chips_per_host (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] chunk_tensor_by_memory() (in module eformer.serialization.utils) clear_once_cache() (eformer.loggings.LazyLogger method) clear_registered_axes() (eformer.escale.partition.manager.PartitionAxis class method) client (eformer.paths.GCSPath attribute) clip_grad (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] clipping_threshold (eformer.optimizers._config.AdafactorConfig attribute), [1] close() (eformer.executor.ray.executor.TpuRemoteManager method) ClusterManager (class in eformer.executor.patch_tpus_ray) ColorFormatter (class in eformer.executor.patch_tpus_ray) (class in eformer.loggings) ColumnWise (class in eformer.common_types) combine_leaf_predicate() (in module eformer.jaximus._imus) command (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] commute_ops (eformer.ops.quantization.implicit_array_1bit.Array1B attribute) (eformer.ops.quantization.implicit_array_8bit.Array8B attribute) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 attribute) (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary attribute) (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary attribute) complete() (eformer.loggings.ProgressLogger method) CompositeShardingRule (class in eformer.escale.helpers.base) compress_bytes() (in module eformer.pytree._serialization) compute_checksum() (eformer.serialization.async_manager.AsyncCheckpointManager static method) compute_dtype (eformer.mpric.policy.policy.Policy attribute), [1] compute_dtype() (eformer.jaximus._imus.ImplicitArray method) compute_shape() (eformer.jaximus._imus.ImplicitArray method) ComputeResourceConfig (class in eformer.executor.ray.resource_manager) concatenate_1bit_operands() (in module eformer.ops.quantization.implicit_array_1bit) concatenate_8bit_operands() (in module eformer.ops.quantization.implicit_array_8bit) concatenate_nf4_operands() (in module eformer.ops.quantization.implicit_array_nf4) config (eformer.optimizers._base.OptimizerBuilder attribute), [1] (eformer.optimizers._base.SchedulerBuilder attribute), [1] (eformer.optimizers._builders.AdafactorOptimizer attribute), [1] (eformer.optimizers._builders.AdamWOptimizer attribute), [1] (eformer.optimizers._builders.ConstantSchedulerBuilder attribute), [1] (eformer.optimizers._builders.CosineSchedulerBuilder attribute), [1] (eformer.optimizers._builders.LinearSchedulerBuilder attribute), [1] (eformer.optimizers._builders.LionOptimizer attribute), [1] (eformer.optimizers._builders.MarsOptimizer attribute), [1] (eformer.optimizers._builders.MuonOptimizer attribute), [1] (eformer.optimizers._builders.QuadOptimizer attribute), [1] (eformer.optimizers._builders.RMSPropOptimizer attribute), [1] (eformer.optimizers._builders.SkewOptimizer attribute), [1]

configure_gcp_artifact_registry() (in module eformer.executor.docker_tpu) ConstantSchedulerBuilder (class in eformer.optimizers._builders) conv_general_dilated_1bit_lhs_rhs() (in module eformer.ops.quantization.implicit_array_1bit) conv_general_dilated_8bit_lhs_rhs() (in module eformer.ops.quantization.implicit_array_8bit) conv_general_dilated_nf4_lhs_rhs() (in module eformer.ops.quantization.implicit_array_nf4) convert_element_type_1bit_operand_kw() (in module eformer.ops.quantization.implicit_array_1bit) convert_element_type_1bit_operand_pos() (in module eformer.ops.quantization.implicit_array_1bit) convert_element_type_8bit_operand_kw() (in module eformer.ops.quantization.implicit_array_8bit) convert_element_type_8bit_operand_pos() (in module eformer.ops.quantization.implicit_array_8bit) convert_element_type_nf4_operand_kw() (in module eformer.ops.quantization.implicit_array_nf4) convert_element_type_nf4_operand_pos() (in module eformer.ops.quantization.implicit_array_nf4) convert_sharding_strategy() (in module eformer.escale.partition.auto_spec) coord_port (eformer.executor.ray.executor.TpuRemoteManager attribute) coordinator_address (eformer.executor.cluster_util.DistributedConfig attribute), [1] coordinator_ip (eformer.executor.ray.types.MultisliceInfo attribute), [1] copy_path() (in module eformer.executor.docker_tpu) copy_tree() (eformer.paths.MLUtilPath method) core_count (eformer.executor.ray.resource_manager.CpuAcceleratorConfig attribute), [1] CosineSchedulerBuilder (class in eformer.optimizers._builders) count (eformer.optimizers._tx.mars.ScaleByMarsState attribute), [1] (eformer.optimizers._tx.utils.OptaxScheduledWeightDecayState attribute), [1] counter (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale attribute), [1] cpu_context() (in module eformer.escale.mesh.creation) cpu_count (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] CpuAcceleratorConfig (class in eformer.executor.ray.resource_manager) create() (eformer.optimizers._factory.OptimizerFactory class method) (eformer.optimizers._factory.OptimizerFactory method) create_actor() (eformer.executor.ray.pool_manager.ResourcePoolManager method) (eformer.executor.ray.pool_manager.SlicePoolManager method) create_cosine_scheduler() (in module eformer.optimizers._tx.utils) create_cpu_mesh() (in module eformer.escale.mesh.creation) create_docker_run_command() (in module eformer.executor.docker_tpu) create_gcs_client() (eformer.serialization.base_manager.CheckpointManager static method) create_linear_scheduler() (in module eformer.optimizers._tx.utils) create_mesh() (in module eformer.escale.mesh.creation) create_monitored_function() (in module eformer.escale.helpers.base) create_partition_spec() (eformer.escale.mesh.mesh_helpers.MeshPartitionHelper method) create_pattern_based_partition_spec() (in module eformer.escale.partition.constraints) create_remote_decorator() (eformer.executor.ray.resource_manager.ComputeResourceConfig method) create_scheduler() (eformer.optimizers._factory.SchedulerFactory method) (eformer.optimizers._factory.SchedulerFactory static method) create_sharding_tree_from_index() (in module eformer.serialization.sharding_utils) create_step_profiler() (in module eformer.loggings) current_actor_handle() (in module eformer.executor.ray.types) current_path() (in module eformer.pytree._serialization) custom_metadata (eformer.serialization.async_manager.CheckpointMetadata attribute), [1]

D

data_fields (eformer.pytree._pytree.PyTreeClassInfo attribute) data_parallel_axis (eformer.escale.partition.manager.PartitionAxis attribute) dataclass() (in module eformer.pytree._xtree) dataclass_types (eformer.aparser._aparser.DataClassArgumentParser attribute), [1] DataClassArgumentParser (class in eformer.aparser._aparser) debug_once() (eformer.loggings.LazyLogger method) decay (eformer.optimizers._config.RMSPropConfig attribute), [1] decay_offset (eformer.optimizers._config.AdafactorConfig attribute), [1] decay_rate (eformer.optimizers._config.AdafactorConfig attribute), [1] decode_attention_dim_axis (eformer.escale.partition.manager.PartitionAxis attribute) decode_attention_kv_dim_axis (eformer.escale.partition.manager.PartitionAxis attribute) decode_batch_axis (eformer.escale.partition.manager.PartitionAxis attribute) decode_head_axis (eformer.escale.partition.manager.PartitionAxis attribute) decode_key_sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute) decode_kv_head_axis (eformer.escale.partition.manager.PartitionAxis attribute) decode_query_sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute) decompress_bytes() (in module eformer.pytree._serialization) deepcopy_tree() (in module eformer.pytree._tree_util) default_handler() (eformer.jaximus._imus.ImplicitArray class method) (in module eformer.jaximus._imus) DEFAULT_MASK_VALUE (in module eformer.common_types) delete() (eformer.ops.quantization.implicit_array_1bit.Array1B method) (eformer.ops.quantization.implicit_array_8bit.Array8B method) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method) DenseState (class in eformer.optimizers._tx.white_kron) dequantize() (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method) dequantize_int8() (in module eformer.ops.quantization.quantization_functions) dequantize_nf4() (in module eformer.ops.quantization.quantization_functions) derive_base_prefix_from_path() (in module eformer.serialization.utils) deserialize_config() (eformer.optimizers._factory.OptimizerFactory class method) (eformer.optimizers._factory.OptimizerFactory method)

device_count (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] device_ids (eformer.common_types.Mesh property) device_remote() (in module eformer.executor.ray.executor) devices (eformer.common_types.Mesh attribute) diag_left (eformer.optimizers._tx.white_kron.LeafState attribute), [1] diag_right (eformer.optimizers._tx.white_kron.LeafState attribute), [1] DistributedConfig (class in eformer.executor.cluster_util) div_1bit_xy() (in module eformer.ops.quantization.implicit_array_1bit) div_1bit_xy_fallback() (in module eformer.ops.quantization.implicit_array_1bit) div_8bit_xy() (in module eformer.ops.quantization.implicit_array_8bit) div_nf4_xy() (in module eformer.ops.quantization.implicit_array_nf4) DockerConfig (class in eformer.executor.ray.docker_executor) DONE (in module eformer.executor.ray.types) DoneSentinel (class in eformer.executor.ray.types) dot() (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary method) (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary method) dot_general_1bit_lhs_rhs() (in module eformer.ops.quantization.implicit_array_1bit) dot_general_8bit_lhs_rhs() (in module eformer.ops.quantization.implicit_array_8bit) dot_general_nf4_lhs_rhs() (in module eformer.ops.quantization.implicit_array_nf4) drain_actor_pool() (eformer.executor.ray.pool_manager.ResourcePoolManager method) (eformer.executor.ray.pool_manager.SlicePoolManager method) dtype (eformer.jaximus._imus.ImplicitArray attribute), [1] (eformer.ops.quantization._config.QuantizationConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] DType (in module eformer.common_types) DTYPE_MAPPING (in module eformer.mpric.dtypes.precision_types) dtype_momentum (eformer.optimizers._config.AdafactorConfig attribute), [1] DTYPE_TO_STRING_MAP (in module eformer.mpric.dtypes.precision_types) DynamicLossScale (class in eformer.mpric.loss_scaling.loss_scaler) DynamicShardingAxes (class in eformer.common_types)

E

eformer.aparser._aparser module eformer.common_types module eformer.escale.helpers.base module eformer.escale.mesh.creation module eformer.escale.mesh.mesh_helpers module eformer.escale.partition.auto_spec module eformer.escale.partition.constraints module eformer.escale.partition.manager module eformer.executor.cluster_util module eformer.executor.docker_tpu module eformer.executor.patch_tpus_ray module eformer.executor.ray.docker_executor module eformer.executor.ray.executor module eformer.executor.ray.pool_manager module eformer.executor.ray.resource_manager module eformer.executor.ray.types module eformer.jaximus._imus module eformer.loggings module eformer.mpric.dtypes.precision_types module eformer.mpric.handler.precision_handler module eformer.mpric.loss_scaling.loss_scaler module eformer.mpric.policy.policy module eformer.ops.quantization._config module eformer.ops.quantization.implicit_array_1bit module eformer.ops.quantization.implicit_array_8bit module eformer.ops.quantization.implicit_array_nf4 module eformer.ops.quantization.implicit_array_rsr module eformer.ops.quantization.quantization_functions module eformer.optimizers._base module eformer.optimizers._builders module eformer.optimizers._config module eformer.optimizers._factory module eformer.optimizers._stage_local module eformer.optimizers._tx.mars module eformer.optimizers._tx.utils module eformer.optimizers._tx.white_kron module eformer.paths module eformer.pytree._pytree module eformer.pytree._serialization module eformer.pytree._tree_util module eformer.pytree._xtree module

eformer.serialization.async_manager module eformer.serialization.base_manager module eformer.serialization.checkpointer module eformer.serialization.fsspec_utils module eformer.serialization.serialization module eformer.serialization.sharding_utils module eformer.serialization.utils module EMBED (in module eformer.common_types) empty (eformer.common_types.Mesh property) empty_node (in module eformer.pytree._tree_util) EMPTY_VAL (in module eformer.common_types) enable (eformer.serialization.async_manager.AsyncCheckpointManager attribute) (eformer.serialization.base_manager.CheckpointManager attribute) enable_compression (eformer.serialization.async_manager.AsyncCheckpointManager attribute) enable_validation (eformer.serialization.async_manager.AsyncCheckpointManager attribute) ensure() (eformer.executor.ray.executor.TpuRemoteManager method) environment (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] eps (eformer.optimizers._config.AdafactorConfig attribute), [1] (eformer.optimizers._config.AdamWConfig attribute), [1] (eformer.optimizers._config.MuonConfig attribute), [1] (eformer.optimizers._config.RMSPropConfig attribute), [1] eps_root (eformer.optimizers._config.AdamWConfig attribute), [1] epsilon (eformer.optimizers._config.MarsConfig attribute), [1] (eformer.optimizers._config.ScionConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] error (eformer.executor.ray.types.JobError attribute), [1] (eformer.executor.ray.types.JobFailed attribute), [1] (eformer.executor.ray.types.JobPreempted attribute), [1] ERROR (eformer.pytree._tree_util.NonePolicy attribute), [1] error_once() (eformer.loggings.LazyLogger method) eSlurmCluster (class in eformer.executor.cluster_util) estimate_array_nbytes() (in module eformer.serialization.utils) estimate_available_memory() (in module eformer.serialization.utils) estimate_memory_usage() (eformer.escale.helpers.base.ShardingAnalyzer method) every (eformer.serialization.checkpointer.CheckpointInterval attribute), [1] ex (eformer.executor.ray.types.ExceptionInfo attribute), [1] ExceptionInfo (class in eformer.executor.ray.types) execute() (eformer.executor.ray.executor.RayExecutor method) (eformer.executor.ray.executor.RayExecutor static method) (in module eformer.executor.ray.executor) execute_on_each_host() (eformer.executor.ray.pool_manager.SlicePoolManager method) execute_on_each_host_flat() (eformer.executor.ray.pool_manager.SlicePoolManager method) execute_on_each_slice() (eformer.executor.ray.pool_manager.SlicePoolManager method) execute_resumable() (eformer.executor.ray.executor.RayExecutor class method) (eformer.executor.ray.executor.RayExecutor method) (in module eformer.executor.ray.executor) execution_env (eformer.executor.ray.resource_manager.ComputeResourceConfig attribute), [1] (eformer.executor.ray.resource_manager.CpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] exists() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) (in module eformer.serialization.fsspec_utils) exp_1bit_x() (in module eformer.ops.quantization.implicit_array_1bit) exp_8bit_x() (in module eformer.ops.quantization.implicit_array_8bit) exp_nf4_x() (in module eformer.ops.quantization.implicit_array_nf4) expand_glob() (in module eformer.serialization.fsspec_utils) EXPERT (in module eformer.common_types) expert_axis (eformer.escale.partition.manager.PartitionAxis attribute) EXPERT_GATE (in module eformer.common_types) expert_gate_axis (eformer.escale.partition.manager.PartitionAxis attribute) expert_parallel_axis (eformer.escale.partition.manager.PartitionAxis attribute) ExpertActivations (class in eformer.common_types) ExpertActivationsAlt (class in eformer.common_types) ExpertColumnWise (class in eformer.common_types) ExpertColumnWiseAlt (class in eformer.common_types) ExpertRowWise (class in eformer.common_types) ExpertRowWiseAlt (class in eformer.common_types) ExpertTensorParallel (class in eformer.common_types) exponent (eformer.optimizers._config.SchedulerConfig attribute), [1] extinguish_profiler() (in module eformer.loggings) extra_kwargs (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] extract_sharding_structure() (in module eformer.escale.partition.constraints) extract_shardings() (in module eformer.escale.partition.constraints)

F

factor (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale attribute), [1] factored (eformer.optimizers._config.AdafactorConfig attribute), [1] failed (eformer.executor.ray.types.HostInfo attribute), [1] field() (in module eformer.pytree._pytree) (in module eformer.pytree._xtree) find_latest_checkpoint() (in module eformer.serialization.checkpointer) flatten_dict() (in module eformer.pytree._tree_util) flatten_for_broadcast() (in module eformer.serialization.utils) flatten_mapping() (in module eformer.pytree._tree_util) flatten_one_implicit_layer() (in module eformer.jaximus._imus) flatten_to_sequence() (in module eformer.pytree._tree_util) flatten_tree() (in module eformer.pytree._tree_util) float_dtype (eformer.serialization.async_manager.AsyncCheckpointManager attribute) (eformer.serialization.base_manager.CheckpointManager attribute) force_cpu() (in module eformer.escale.mesh.creation) forkify_remote_fn() (eformer.executor.ray.resource_manager.RayResources static method) format() (eformer.executor.patch_tpus_ray.ColorFormatter method) (eformer.loggings.ColorFormatter method) framework_version (eformer.serialization.async_manager.CheckpointMetadata attribute), [1] from_bytes() (in module eformer.pytree._serialization) from_compressed_bytes() (in module eformer.pytree._serialization) from_dict() (eformer.escale.partition.manager.PartitionManager class method) (eformer.optimizers._config.SerializationMixin class method) (eformer.optimizers._config.SerializationMixin method) (eformer.optimizers._tx.white_kron.DenseState class method) (eformer.optimizers._tx.white_kron.LeafState class method) (eformer.pytree._tree_util.StateValidationResult class method) (eformer.pytree._xtree.xTree class method) (eformer.serialization.async_manager.CheckpointMetadata class method)

from_json() (eformer.escale.partition.manager.PartitionManager class method) (eformer.optimizers._config.SerializationMixin class method) (eformer.optimizers._config.SerializationMixin method) (eformer.optimizers._tx.white_kron.DenseState class method) (eformer.optimizers._tx.white_kron.LeafState class method) (eformer.pytree._tree_util.StateValidationResult class method) from_matrix() (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary class method) (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary class method) from_resource_dict() (eformer.executor.ray.resource_manager.RayResources static method) from_state_dict() (eformer.pytree._xtree.xTree class method) (in module eformer.pytree._serialization) from_string() (eformer.mpric.policy.policy.Policy class method) frozen (eformer.pytree._pytree.PyTreeClassInfo attribute) FrozenPyTree (class in eformer.pytree._pytree) fully_sharded_data_parallel_axis (eformer.escale.partition.manager.PartitionAxis attribute)

G

gamma (eformer.optimizers._config.MarsConfig attribute), [1] gather_1bit_operand() (in module eformer.ops.quantization.implicit_array_1bit) gather_8bit_operand() (in module eformer.ops.quantization.implicit_array_8bit) gather_nf4_operand() (in module eformer.ops.quantization.implicit_array_nf4) gcs_bucket (eformer.serialization.async_manager.AsyncCheckpointManager attribute) (eformer.serialization.base_manager.CheckpointManager attribute) gcs_client (eformer.paths.PathManager attribute) (eformer.paths.PathManager property) (eformer.serialization.base_manager.CheckpointManager attribute) GCSPath (class in eformer.paths) generate_template() (eformer.optimizers._factory.OptimizerFactory class method) (eformer.optimizers._factory.OptimizerFactory method) GENERATION_MODES (in module eformer.common_types) get() (eformer.executor.ray.types.RefBox method) get_actor_name_from_actor_info() (eformer.executor.ray.pool_manager.ResourcePoolManager method) (eformer.executor.ray.pool_manager.SlicePoolManager method) get_actor_pool_name() (eformer.executor.ray.pool_manager.ResourcePoolManager method) (eformer.executor.ray.pool_manager.SlicePoolManager method) get_all_actors_in_pool() (eformer.executor.ray.pool_manager.ResourcePoolManager method) get_all_pool_members() (eformer.executor.ray.pool_manager.ResourcePoolManager method) get_axes_size_in_mesh() (in module eformer.escale.partition.constraints) get_base_optimizer() (in module eformer.optimizers._tx.utils) get_coordinator_address() (eformer.executor.cluster_util.eSlurmCluster class method) get_corrected_named_sharding() (in module eformer.escale.partition.constraints) get_count() (eformer.pytree._tree_util.MetaValueRecreator method) get_current_partition_manager() (in module eformer.escale.partition.manager) get_external_ip() (eformer.executor.patch_tpus_ray.IPManager static method) get_incontext_mesh() (in module eformer.escale.partition.constraints) get_local_device_ids_for_process() (eformer.executor.cluster_util.eSlurmCluster class method) get_local_ip() (eformer.executor.patch_tpus_ray.IPManager static method) get_logger() (in module eformer.executor.patch_tpus_ray) (in module eformer.loggings) get_mesh_axis_names() (in module eformer.escale.partition.constraints)

get_mesh_axis_size() (in module eformer.escale.partition.constraints) get_names_from_partition_spec() (in module eformer.escale.partition.constraints) get_nf4() (in module eformer.ops.quantization.quantization_functions) get_opt_state_partition_specs() (in module eformer.optimizers._tx.white_kron) get_partition_manager() (in module eformer.escale.partition.manager) get_partition_spec() (in module eformer.escale.partition.constraints) get_platform_default_half() (in module eformer.mpric.dtypes.precision_types) get_registered_axes() (eformer.escale.partition.manager.PartitionAxis class method) get_remote_options() (eformer.executor.ray.resource_manager.ComputeResourceConfig method) (eformer.executor.ray.resource_manager.CpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.GpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.TpuAcceleratorConfig method) get_rng() (eformer.pytree._tree_util.MetaValueRecreator method) get_safe_hash_int() (in module eformer.escale.partition.manager) get_serialization_info() (in module eformer.pytree._serialization) get_shardings_with_structure() (in module eformer.escale.partition.constraints) get_submesh_device_index() (in module eformer.escale.partition.constraints) glob() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) global_manager (eformer.serialization.async_manager.AsyncCheckpointManager property) GOOGLE_TPU_V2 (eformer.executor.ray.resource_manager.HardwareType attribute) GOOGLE_TPU_V3 (eformer.executor.ray.resource_manager.HardwareType attribute) GOOGLE_TPU_V4 (eformer.executor.ray.resource_manager.HardwareType attribute) GOOGLE_TPU_V5LITEPOD (eformer.executor.ray.resource_manager.HardwareType attribute) GOOGLE_TPU_V5P (eformer.executor.ray.resource_manager.HardwareType attribute) GOOGLE_TPU_V6E (eformer.executor.ray.resource_manager.HardwareType attribute) gpu_model (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] GpuAcceleratorConfig (class in eformer.executor.ray.resource_manager) gpus (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] gradient_accumulation_steps (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] group_keys_by_shard_size() (in module eformer.serialization.utils) growth_interval (eformer.mpric.loss_scaling.loss_scaler.LossScaleConfig attribute), [1]

H

handle_ray_error() (in module eformer.executor.ray.types) haps (eformer.optimizers._config.SoapConfig attribute), [1] hardware_identifier() (eformer.executor.ray.resource_manager.ComputeResourceConfig method) (eformer.executor.ray.resource_manager.CpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.GpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.TpuAcceleratorConfig method) HardwareType (class in eformer.executor.ray.resource_manager) hash_fn() (in module eformer.escale.partition.manager) head_axis (eformer.escale.partition.manager.PartitionAxis attribute)

head_name (eformer.executor.ray.resource_manager.ComputeResourceConfig attribute), [1] head_workers (eformer.executor.ray.resource_manager.ComputeResourceConfig attribute), [1] healthy (eformer.executor.ray.types.HostInfo attribute), [1] hidden_state_axis (eformer.escale.partition.manager.PartitionAxis attribute) HiddenStateSharding (class in eformer.common_types) host_id (eformer.executor.ray.types.HostInfo attribute), [1] host_infos (eformer.executor.ray.types.SliceInfo attribute) HostInfo (class in eformer.executor.ray.types) HUAWEI_NPU_910B (eformer.executor.ray.resource_manager.HardwareType attribute) HUAWEI_NPU_910B4 (eformer.executor.ray.resource_manager.HardwareType attribute)

I

i4tou4() (in module eformer.ops.quantization.quantization_functions) i8tou8() (in module eformer.ops.quantization.quantization_functions) ignite_profiler() (in module eformer.loggings) image (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] implicit() (in module eformer.jaximus._imus) implicit_depth() (in module eformer.jaximus._imus) ImplicitArray (class in eformer.jaximus._imus) index_filename() (in module eformer.serialization.utils) inference_wrapper() (eformer.mpric.handler.precision_handler.PrecisionHandler method) info (eformer.executor.ray.types.JobStatus attribute), [1] info_once() (eformer.loggings.LazyLogger method) initial_scale (eformer.mpric.loss_scaling.loss_scaler.LossScaleConfig attribute), [1] (eformer.optimizers._config.RMSPropConfig attribute), [1] initialize() (eformer.executor.cluster_util.DistributedConfig method) (eformer.executor.cluster_util.RayClusterConfig method) InsufficientSlicesError INT8 (eformer.ops.quantization._config.QuantizationType attribute) int_key_to_string() (in module eformer.pytree._tree_util) integer_pow_1bit_x() (in module eformer.ops.quantization.implicit_array_1bit) integer_pow_1bit_xy() (in module eformer.ops.quantization.implicit_array_1bit) integer_pow_8bit_x() (in module eformer.ops.quantization.implicit_array_8bit) integer_pow_8bit_xy() (in module eformer.ops.quantization.implicit_array_8bit) integer_pow_nf4_x() (in module eformer.ops.quantization.implicit_array_nf4) integer_pow_nf4_xy() (in module eformer.ops.quantization.implicit_array_nf4) INTEL_GAUDI (eformer.executor.ray.resource_manager.HardwareType attribute) INTEL_MAX_1100 (eformer.executor.ray.resource_manager.HardwareType attribute) INTEL_MAX_1550 (eformer.executor.ray.resource_manager.HardwareType attribute) invalid_types (eformer.pytree._tree_util.StateValidationResult attribute), [1] ip_address (eformer.executor.ray.types.SliceInfo attribute), [1] IPManager (class in eformer.executor.patch_tpus_ray)

is_absolute() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) is_array() (in module eformer.jaximus._imus) (in module eformer.pytree._tree_util) is_array_like() (in module eformer.pytree._tree_util) (in module eformer.serialization.serialization) is_dir() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) is_file() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) is_flatten() (in module eformer.pytree._tree_util) is_gcs_path() (in module eformer.serialization.utils) is_iterable() (in module eformer.pytree._tree_util) is_kernel_available() (in module eformer.ops.quantization.quantization_functions) is_local_path() (in module eformer.paths) is_multi_process (eformer.common_types.Mesh property) is_ray_running() (eformer.executor.patch_tpus_ray.RayManager method) is_remote_path() (in module eformer.paths) is_scalar (eformer.common_types.Mesh property) is_serializable() (in module eformer.pytree._serialization) is_sharded (eformer.ops.quantization.implicit_array_8bit.Array8B property) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 property) is_tensorstore() (eformer.serialization.async_manager.AsyncCheckpointManager static method) is_valid (eformer.pytree._tree_util.StateValidationResult attribute), [1] iterdir() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method)

J

JobError (class in eformer.executor.ray.types) JobFailed (class in eformer.executor.ray.types) JobInfo (class in eformer.executor.ray.types) JobPreempted (class in eformer.executor.ray.types)

JobStatus (class in eformer.executor.ray.types) JobSucceeded (class in eformer.executor.ray.types) join_key() (in module eformer.pytree._tree_util) join_path() (in module eformer.serialization.fsspec_utils)

K

k (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary attribute), [1] key_path_to_str() (in module eformer.pytree._tree_util) key_sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute) kind (eformer.executor.ray.types.JobInfo attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1]

KronConfig (class in eformer.optimizers._config) KV_HEAD (in module eformer.common_types) kv_head_axis (eformer.escale.partition.manager.PartitionAxis attribute) KV_HEAD_DIM (in module eformer.common_types) KV_LENGTH (in module eformer.common_types)

L

lax_map_batch_size (eformer.optimizers._config.KronConfig attribute), [1] lax_map_scanned_layers (eformer.optimizers._config.KronConfig attribute), [1] LazyLogger (class in eformer.executor.patch_tpus_ray) (class in eformer.loggings) leaf_key_paths() (in module eformer.pytree._tree_util) (in module eformer.serialization.serialization) leaf_predicate() (in module eformer.jaximus._imus) LeafState (class in eformer.optimizers._tx.white_kron) learning_rate (eformer.optimizers._config.SchedulerConfig attribute), [1] learning_rate_end (eformer.optimizers._config.SchedulerConfig attribute), [1] LENGTH (in module eformer.common_types) level (eformer.loggings.LazyLogger attribute) (eformer.loggings.LazyLogger property) LinearSchedulerBuilder (class in eformer.optimizers._builders) LionConfig (class in eformer.optimizers._config) LionOptimizer (class in eformer.optimizers._builders) Ll (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] load() (eformer.serialization.async_manager.AsyncCheckpointManager method) load_checkpoint() (eformer.serialization.base_manager.CheckpointManager static method) (eformer.serialization.checkpointer.Checkpointer method) load_dict() (eformer.paths.MLUtilPath method) load_from_file() (in module eformer.pytree._serialization) load_jax_array() (eformer.paths.MLUtilPath method) load_pytree() (eformer.serialization.async_manager.AsyncCheckpointManager method) (eformer.serialization.checkpointer.Checkpointer method)

load_state_from_gcs_msgpack() (eformer.serialization.base_manager.CheckpointManager method) load_tree_parallel() (eformer.serialization.async_manager.AsyncCheckpointManager method) local_device_ids (eformer.executor.cluster_util.DistributedConfig attribute), [1] local_devices (eformer.common_types.Mesh property) local_mesh (eformer.common_types.Mesh property) LocalPath (class in eformer.paths) log_1bit_x() (in module eformer.ops.quantization.implicit_array_1bit) log_8bit_x() (in module eformer.ops.quantization.implicit_array_8bit) log_failures_to() (in module eformer.executor.ray.types) log_nf4_x() (in module eformer.ops.quantization.implicit_array_nf4) logger (eformer.executor.patch_tpus_ray.ClusterManager attribute) (eformer.executor.patch_tpus_ray.RayManager attribute) (eformer.executor.ray.types.SnitchRecipient attribute), [1] logical_cpu_core_count() (in module eformer.executor.cluster_util) loss_scale (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale attribute), [1] (eformer.mpric.loss_scaling.loss_scaler.NoOpLossScale attribute) (eformer.mpric.loss_scaling.loss_scaler.NoOpLossScale property) loss_scale_config (eformer.mpric.handler.precision_handler.PrecisionHandler attribute) loss_scaler (eformer.mpric.handler.precision_handler.PrecisionHandler attribute) LossScaleConfig (class in eformer.mpric.loss_scaling.loss_scaler) Lr (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] lr_style (eformer.optimizers._config.WhiteKronConfig attribute), [1] lt_1bit_xy() (in module eformer.ops.quantization.implicit_array_1bit) lt_8bit_xy() (in module eformer.ops.quantization.implicit_array_8bit) lt_nf4_xy() (in module eformer.ops.quantization.implicit_array_nf4)

M

main() (in module eformer.executor.patch_tpus_ray) make_choice_type_function() (in module eformer.aparser._aparser) make_docker_run_command() (in module eformer.executor.ray.docker_executor) make_itsharded() (in module eformer.serialization.sharding_utils) make_shard_and_gather_fns() (in module eformer.escale.partition.constraints) make_stage_local_gradient_transformation() (in module eformer.optimizers._stage_local) make_unsupported_stage_local_gradient_transformation() (in module eformer.optimizers._stage_local) manage_extra_context() (in module eformer.executor.docker_tpu) mars() (in module eformer.optimizers._tx.mars) MarsConfig (class in eformer.optimizers._config) MarsOptimizer (class in eformer.optimizers._builders) match_partition_rules() (in module eformer.escale.partition.constraints) materialize() (eformer.jaximus._imus.ImplicitArray method) (eformer.ops.quantization.implicit_array_1bit.Array1B method), [1] (eformer.ops.quantization.implicit_array_8bit.Array8B method), [1] (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method), [1] (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary method) (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary method) materialize_handler() (in module eformer.jaximus._imus) materialize_nested() (in module eformer.jaximus._imus) matmul_bf16_int8_weight_only() (in module eformer.ops.quantization.implicit_array_8bit) max_1bit_xy() (in module eformer.ops.quantization.implicit_array_1bit) max_8bit_xy() (in module eformer.ops.quantization.implicit_array_8bit) max_grad_norm (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.MarsConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] max_memory_per_device (eformer.escale.helpers.base.MemoryConstrainedShardingRule attribute) max_nf4_xy() (in module eformer.ops.quantization.implicit_array_nf4) max_precond_dim (eformer.optimizers._config.SoapConfig attribute), [1] max_size_dense (eformer.optimizers._config.WhiteKronConfig attribute), [1] max_size_triangular (eformer.optimizers._config.KronConfig attribute), [1] memory_save_mode (eformer.optimizers._config.KronConfig attribute), [1] MemoryConstrainedShardingRule (class in eformer.escale.helpers.base) merge() (in module eformer.pytree._tree_util) merge_small_dims (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] merged (eformer.optimizers._tx.white_kron.LeafState attribute), [1] Mesh (class in eformer.common_types) mesh (eformer.escale.helpers.base.AutoShardingRule attribute) (eformer.escale.helpers.base.MemoryConstrainedShardingRule attribute) (eformer.escale.helpers.base.ShardingAnalyzer attribute) (eformer.escale.mesh.mesh_helpers.MeshPartitionHelper attribute) MeshPartitionHelper (class in eformer.escale.mesh.mesh_helpers) meta_fields (eformer.pytree._pytree.PyTreeClassInfo attribute) MetaValueRecreator (class in eformer.pytree._tree_util) min_dim_size_to_factor (eformer.optimizers._config.AdafactorConfig attribute), [1] min_loss_scale (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale attribute), [1] min_ndim_triangular (eformer.optimizers._config.KronConfig attribute), [1] min_scale (eformer.mpric.loss_scaling.loss_scaler.LossScaleConfig attribute), [1] min_shard_size (eformer.escale.helpers.base.AutoShardingRule attribute) missing_keys (eformer.pytree._tree_util.StateValidationResult attribute), [1] mkdir() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) mkdirs() (in module eformer.serialization.fsspec_utils) MLP_INTERMEDIATE (in module eformer.common_types) mlp_intermediate_axis (eformer.escale.partition.manager.PartitionAxis attribute) MLUtilPath (class in eformer.paths) mode (eformer.common_types.AttnKVSharding attribute) (eformer.common_types.AttnQSharding attribute) (eformer.common_types.ColumnWise attribute) (eformer.common_types.DynamicShardingAxes attribute), [1] (eformer.common_types.ExpertActivations attribute) (eformer.common_types.ExpertActivationsAlt attribute) (eformer.common_types.ExpertColumnWise attribute) (eformer.common_types.ExpertColumnWiseAlt attribute) (eformer.common_types.ExpertRowWise attribute) (eformer.common_types.ExpertRowWiseAlt attribute) (eformer.common_types.ExpertTensorParallel attribute) (eformer.common_types.HiddenStateSharding attribute) (eformer.common_types.Replicated attribute) (eformer.common_types.RowWise attribute) (eformer.common_types.SColumnWise attribute) (eformer.common_types.SRowWise attribute) (eformer.common_types.UnifiedExpertColumnWise attribute) (eformer.common_types.UnifiedExpertRowWise attribute)

MODE_DECODE (in module eformer.common_types) MODE_INSERT (in module eformer.common_types) MODE_PREFILL (in module eformer.common_types) MODE_TRAIN (in module eformer.common_types) module eformer.aparser._aparser eformer.common_types eformer.escale.helpers.base eformer.escale.mesh.creation eformer.escale.mesh.mesh_helpers eformer.escale.partition.auto_spec eformer.escale.partition.constraints eformer.escale.partition.manager eformer.executor.cluster_util eformer.executor.docker_tpu eformer.executor.patch_tpus_ray eformer.executor.ray.docker_executor eformer.executor.ray.executor eformer.executor.ray.pool_manager eformer.executor.ray.resource_manager eformer.executor.ray.types eformer.jaximus._imus eformer.loggings eformer.mpric.dtypes.precision_types eformer.mpric.handler.precision_handler eformer.mpric.loss_scaling.loss_scaler eformer.mpric.policy.policy eformer.ops.quantization._config eformer.ops.quantization.implicit_array_1bit eformer.ops.quantization.implicit_array_8bit eformer.ops.quantization.implicit_array_nf4 eformer.ops.quantization.implicit_array_rsr eformer.ops.quantization.quantization_functions eformer.optimizers._base eformer.optimizers._builders eformer.optimizers._config eformer.optimizers._factory eformer.optimizers._stage_local eformer.optimizers._tx.mars eformer.optimizers._tx.utils eformer.optimizers._tx.white_kron eformer.paths eformer.pytree._pytree eformer.pytree._serialization eformer.pytree._tree_util eformer.pytree._xtree eformer.serialization.async_manager eformer.serialization.base_manager eformer.serialization.checkpointer eformer.serialization.fsspec_utils eformer.serialization.serialization eformer.serialization.sharding_utils eformer.serialization.utils mog (eformer.optimizers._tx.mars.ScaleByMarsState attribute), [1] momentum (eformer.optimizers._config.AdafactorConfig attribute), [1] (eformer.optimizers._config.RMSPropConfig attribute), [1] (eformer.optimizers._config.ScionConfig attribute), [1] msgpack_restore() (in module eformer.pytree._serialization) msgpack_serialize() (in module eformer.pytree._serialization) mu (eformer.optimizers._tx.mars.ScaleByMarsState attribute), [1] mu_dtype (eformer.optimizers._config.AdamWConfig attribute), [1] (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.LionConfig attribute), [1] (eformer.optimizers._config.MuonConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] mul_1bit_xy() (in module eformer.ops.quantization.implicit_array_1bit) mul_8bit_xy() (in module eformer.ops.quantization.implicit_array_8bit) mul_nf4_xy() (in module eformer.ops.quantization.implicit_array_nf4) multiply_by_parameter_scale (eformer.optimizers._config.AdafactorConfig attribute), [1] MultisliceInfo (class in eformer.executor.ray.types) MuonConfig (class in eformer.optimizers._config) MuonOptimizer (class in eformer.optimizers._builders)

N

name (eformer.executor.ray.types.JobInfo attribute), [1] (eformer.loggings.LazyLogger attribute) (eformer.loggings.LazyLogger property) (eformer.loggings.ProgressLogger attribute) (eformer.paths.GCSPath property) (eformer.paths.LocalPath property) named_tree_map() (in module eformer.pytree._tree_util) names_in_current_mesh() (in module eformer.escale.partition.constraints) nc (eformer.optimizers._tx.white_kron.LeafState attribute), [1] nesterov (eformer.optimizers._config.MuonConfig attribute), [1] (eformer.optimizers._config.RMSPropConfig attribute), [1] network (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] NF4 (eformer.ops.quantization._config.QuantizationType attribute) nf4_matmul() (in module eformer.ops.quantization.quantization_functions) nf4_use_kernel() (in module eformer.ops.quantization.quantization_functions) nf4xf32_to_f32() (in module eformer.ops.quantization.quantization_functions) node_id (eformer.executor.ray.types.HostInfo attribute) node_ids (eformer.executor.ray.types.SliceInfo attribute) noise_scale (eformer.optimizers._config.WhiteKronConfig attribute), [1] NonePolicy (class in eformer.pytree._tree_util) NoOpLossScale (class in eformer.mpric.loss_scaling.loss_scaler) normalize_grads (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] nr (eformer.optimizers._tx.white_kron.LeafState attribute), [1]

ns_coeffs (eformer.optimizers._config.MuonConfig attribute), [1] ns_steps (eformer.optimizers._config.MuonConfig attribute), [1] nu (eformer.optimizers._tx.mars.ScaleByMarsState attribute), [1] num_accelerators_per_host (eformer.executor.ray.types.SliceInfo attribute), [1] num_cpus (eformer.executor.ray.resource_manager.RayResources attribute), [1] num_devices (eformer.executor.ray.types.HostInfo attribute), [1] num_gpus (eformer.executor.ray.resource_manager.RayResources attribute), [1] num_hosts (eformer.executor.ray.types.SliceInfo attribute), [1] num_processes (eformer.executor.cluster_util.DistributedConfig attribute), [1] num_slices (eformer.executor.ray.types.MultisliceInfo attribute), [1] NVIDIA_A100 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_A100_40G (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_A100_80G (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_H100 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_H20 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_H200 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_L4 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_L40S (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_TESLA_A10G (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_TESLA_K80 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_TESLA_P100 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_TESLA_P4 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_TESLA_T4 (eformer.executor.ray.resource_manager.HardwareType attribute) NVIDIA_TESLA_V100 (eformer.executor.ray.resource_manager.HardwareType attribute)

O

offset (eformer.pytree._tree_util.PackedLeaf attribute), [1] on_step() (eformer.serialization.checkpointer.Checkpointer method) one_diag (eformer.optimizers._config.SoapConfig attribute), [1] one_hot_maps (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary attribute), [1] optax_add_scheduled_weight_decay() (in module eformer.optimizers._tx.utils) OptaxScheduledWeightDecayState (class in eformer.optimizers._tx.utils) optimize_shard_layout() (in module eformer.serialization.utils)

optimize_sharding_for_memory() (in module eformer.escale.partition.auto_spec) optimizer_config (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] OptimizerBuilder (class in eformer.optimizers._base) OptimizerFactory (class in eformer.optimizers._factory) org_dtype (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary attribute), [1] OrginArray (class in eformer.jaximus._imus) output_dtype (eformer.mpric.policy.policy.Policy attribute), [1]

P

pack_pytree() (in module eformer.pytree._tree_util) pack_weights_1bit() (in module eformer.ops.quantization.quantization_functions) packed (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 attribute), [1] PackedLeaf (class in eformer.pytree._tree_util) padding (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary attribute), [1] param_dtype (eformer.mpric.policy.policy.Policy attribute), [1] params_partition_specs (eformer.optimizers._config.WhiteKronConfig attribute), [1] parent (eformer.paths.GCSPath property) (eformer.paths.LocalPath property) parse_args_into_dataclasses() (eformer.aparser._aparser.DataClassArgumentParser method) parse_arguments() (eformer.executor.patch_tpus_ray.ArgumentParser static method) parse_dict() (eformer.aparser._aparser.DataClassArgumentParser method) parse_gcs_path() (in module eformer.serialization.utils) parse_image_name_and_tag() (in module eformer.executor.docker_tpu) parse_json_file() (eformer.aparser._aparser.DataClassArgumentParser method) parse_mesh_from_string() (in module eformer.escale.mesh.creation) parse_yaml_file() (eformer.aparser._aparser.DataClassArgumentParser method) partition_grads_into_blocks (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] PartitionAxis (class in eformer.escale.partition.manager) PartitionManager (class in eformer.escale.partition.manager) parts() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) path (eformer.paths.GCSPath attribute) (eformer.paths.LocalPath attribute) path_protocol() (in module eformer.paths) PathManager (class in eformer.paths) paxis (eformer.escale.partition.manager.PartitionManager attribute)

period (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale attribute), [1] pipeline_axis_name (eformer.optimizers._config.WhiteKronConfig attribute), [1] pipeline_axis_size (eformer.optimizers._config.WhiteKronConfig attribute), [1] pod_count (eformer.executor.ray.executor.TpuRemoteManager attribute) (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] Policy (class in eformer.mpric.policy.policy) policy (eformer.mpric.handler.precision_handler.PrecisionHandler attribute) port (eformer.executor.ray.types.MultisliceInfo attribute), [1] PrecisionHandler (class in eformer.mpric.handler.precision_handler) precond_dtype (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] precond_grads_precision (eformer.optimizers._config.KronConfig attribute), [1] precond_update_precision (eformer.optimizers._config.KronConfig attribute), [1] precondition_frequency (eformer.optimizers._config.SoapConfig attribute), [1] preconditioner_init_scale (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] preconditioner_lr (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] preconditioner_update_probability (eformer.optimizers._config.KronConfig attribute), [1] prepare_all_slices() (eformer.executor.ray.pool_manager.SlicePoolManager method) PRESERVE (eformer.pytree._tree_util.NonePolicy attribute), [1] print_remote_raise() (in module eformer.executor.ray.types) privileged (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] PRNGKey (in module eformer.common_types) process_id (eformer.executor.cluster_util.DistributedConfig attribute), [1] ProgressLogger (class in eformer.loggings) push_image_to_gcp() (in module eformer.executor.docker_tpu) push_image_to_github() (in module eformer.executor.docker_tpu) PyTree (class in eformer.pytree._pytree) PyTreeClassInfo (class in eformer.pytree._pytree)

Q

Ql (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] Qr (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] quad() (in module eformer.optimizers._tx.white_kron) QuadOptimizer (class in eformer.optimizers._builders) QuantizationConfig (class in eformer.ops.quantization._config) QuantizationType (class in eformer.ops.quantization._config) quantize() (eformer.ops.quantization.implicit_array_1bit.Array1B class method) (eformer.ops.quantization.implicit_array_8bit.Array8B class method) (eformer.ops.quantization.implicit_array_8bit.Array8B method) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 class method) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method) (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary class method) (in module eformer.ops.quantization._config)

quantize_and_pack_nf4() (in module eformer.ops.quantization.quantization_functions) quantize_int8() (in module eformer.ops.quantization.quantization_functions) QUERY_LENGTH (in module eformer.common_types) query_sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute)

R

ray_manager (eformer.executor.patch_tpus_ray.ClusterManager attribute) ray_path (eformer.executor.patch_tpus_ray.RayManager attribute) RayClusterConfig (class in eformer.executor.cluster_util) RayExecutor (class in eformer.executor.ray.executor) RayManager (class in eformer.executor.patch_tpus_ray) RayResources (class in eformer.executor.ray.resource_manager) read_bytes() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) read_ips_from_yaml() (eformer.executor.patch_tpus_ray.IPManager static method) read_process_array() (in module eformer.serialization.utils) read_text() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) recursive_merge() (in module eformer.pytree._tree_util) redecorate_remote_fn_for_call() (eformer.executor.ray.resource_manager.CpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.GpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.TpuAcceleratorConfig method) reduce_1bit_operand_init_value() (in module eformer.ops.quantization.implicit_array_1bit) reduce_8bit_operand_init_value() (in module eformer.ops.quantization.implicit_array_8bit) reduce_nf4_operand_init_value() (in module eformer.ops.quantization.implicit_array_nf4) ref (eformer.executor.ray.types.RefBox attribute), [1] RefBox (class in eformer.executor.ray.types) register() (eformer.escale.partition.manager.PartitionAxis class method) (in module eformer.jaximus._imus) register_optimizer() (in module eformer.optimizers._base) register_scheduler() (in module eformer.optimizers._base) register_serialization_state() (in module eformer.pytree._serialization) (in module eformer.pytree._xtree) relative_to() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) remove (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] remove() (in module eformer.serialization.fsspec_utils) remove_path() (in module eformer.executor.docker_tpu) rename() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) REPLACE (eformer.pytree._tree_util.NonePolicy attribute), [1] replace() (eformer.escale.partition.manager.PartitionAxis method) (eformer.escale.partition.manager.PartitionManager method) (eformer.optimizers._tx.white_kron.DenseState method) (eformer.optimizers._tx.white_kron.LeafState method) (eformer.pytree._tree_util.StateValidationResult method) (eformer.pytree._xtree.xTree method)

Replicated (class in eformer.common_types) reraise() (eformer.executor.ray.types.ExceptionInfo method) reshape_1bit_operand() (in module eformer.ops.quantization.implicit_array_1bit) reshape_8bit_operand() (in module eformer.ops.quantization.implicit_array_8bit) reshape_nf4_operand() (in module eformer.ops.quantization.implicit_array_nf4) reshard() (eformer.ops.quantization.implicit_array_8bit.Array8B method) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method) (in module eformer.serialization.utils) resolve() (eformer.escale.partition.manager.PartitionManager method) (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) resolve_axis() (eformer.escale.partition.manager.PartitionAxis method) resolve_maybe_refs() (in module eformer.executor.ray.executor) resolve_spec() (eformer.escale.partition.manager.PartitionAxis method) resource_name (eformer.executor.ray.resource_manager.CpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] ResourcePoolManager (class in eformer.executor.ray.pool_manager) resources (eformer.executor.ray.resource_manager.RayResources attribute), [1] restore() (eformer.executor.ray.types.ExceptionInfo method) result (eformer.executor.ray.types.JobSucceeded attribute), [1] reverse (eformer.escale.helpers.base.AutoShardingRule attribute) rmdir() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) RMSPropConfig (class in eformer.optimizers._config) RMSPropOptimizer (class in eformer.optimizers._builders) RowWise (class in eformer.common_types) rsr_b1 (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary attribute), [1] rsr_b2 (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary attribute), [1] RSROperatorBinary (class in eformer.ops.quantization.implicit_array_rsr) RSROperatorTernary (class in eformer.ops.quantization.implicit_array_rsr) rules (eformer.escale.helpers.base.CompositeShardingRule attribute) run_class_method() (eformer.executor.ray.executor.TpuRemoteManager method) run_command() (eformer.executor.patch_tpus_ray.RayManager method) (in module eformer.executor.docker_tpu) run_docker_multislice() (in module eformer.executor.ray.docker_executor) run_docker_on_pod() (in module eformer.executor.ray.docker_executor) run_function() (eformer.executor.ray.executor.TpuRemoteManager method) runtime_env (eformer.executor.ray.executor.TpuRemoteManager attribute) (eformer.executor.ray.resource_manager.RayResources attribute), [1] RUNTIME_MODE_TYPES (in module eformer.common_types) runtime_name (eformer.executor.ray.resource_manager.CpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1]

S

safe_delete() (in module eformer.ops.quantization.implicit_array_nf4) safe_loadpath() (eformer.serialization.async_manager.AsyncCheckpointManager static method) safe_materialize() (in module eformer.ops.quantization.implicit_array_nf4) save() (eformer.serialization.async_manager.AsyncCheckpointManager method) save_checkpoint() (eformer.serialization.base_manager.CheckpointManager class method) (eformer.serialization.checkpointer.Checkpointer method) save_dict() (eformer.paths.MLUtilPath method) save_interval (eformer.serialization.checkpointer.Checkpointer attribute) save_jax_array() (eformer.paths.MLUtilPath method) save_optimizer_state (eformer.serialization.base_manager.CheckpointManager attribute) save_pytree() (eformer.serialization.async_manager.AsyncCheckpointManager method) (eformer.serialization.checkpointer.Checkpointer method) save_state_to_gcs_msgpack() (eformer.serialization.base_manager.CheckpointManager method) save_to_file() (in module eformer.pytree._serialization) save_tree() (eformer.serialization.async_manager.AsyncCheckpointManager method) scale (eformer.ops.quantization.implicit_array_8bit.Array8B attribute), [1] scale() (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale method) (eformer.mpric.loss_scaling.loss_scaler.NoOpLossScale method) scale_by_mars() (in module eformer.optimizers._tx.mars) scale_by_quad() (in module eformer.optimizers._tx.white_kron) scale_by_skew() (in module eformer.optimizers._tx.white_kron) scale_factor (eformer.mpric.loss_scaling.loss_scaler.LossScaleConfig attribute), [1] scale_multislice() (eformer.executor.ray.pool_manager.SlicePoolManager method) ScaleByMarsState (class in eformer.optimizers._tx.mars) scanned (eformer.optimizers._tx.white_kron.LeafState attribute), [1] scanned_layers (eformer.optimizers._config.WhiteKronConfig attribute), [1] schedule_list (eformer.optimizers._config.SoapConfig attribute), [1] schedule_on_each_host() (eformer.executor.ray.pool_manager.SlicePoolManager method) scheduler (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] scheduler_type (eformer.optimizers._config.SchedulerConfig attribute), [1] SchedulerBuilder (class in eformer.optimizers._base) SchedulerConfig (class in eformer.optimizers._config) SchedulerFactory (class in eformer.optimizers._factory) ScionConfig (class in eformer.optimizers._config) SColumnWise (class in eformer.common_types) separate_process_fn() (eformer.executor.ray.resource_manager.RayResources static method) sequence_axis (eformer.escale.partition.manager.PartitionAxis attribute) sequence_parallel_axis (eformer.escale.partition.manager.PartitionAxis attribute) ser_exc_info() (eformer.executor.ray.types.ExceptionInfo class method) SerializationMixin (class in eformer.optimizers._config) serialize_config() (eformer.optimizers._factory.OptimizerFactory class method) (eformer.optimizers._factory.OptimizerFactory method) setup_cluster() (eformer.executor.patch_tpus_ray.ClusterManager method) setup_self_job_node() (eformer.executor.patch_tpus_ray.ClusterManager method) shampoo_beta (eformer.optimizers._config.SoapConfig attribute), [1] shape (eformer.common_types.Mesh property) (eformer.jaximus._imus.ImplicitArray attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] (eformer.pytree._tree_util.PackedLeaf attribute), [1] Shape (in module eformer.common_types) shape_patterns (eformer.escale.helpers.base.ShapeBasedShardingRule attribute) shape_tuple (eformer.common_types.Mesh property) ShapeBasedShardingRule (class in eformer.escale.helpers.base) shard() (eformer.escale.partition.manager.PartitionManager method) shard_array() (eformer.escale.mesh.mesh_helpers.MeshPartitionHelper method) shard_filename() (in module eformer.serialization.utils) sharding (eformer.ops.quantization.implicit_array_8bit.Array8B attribute), [1] (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 attribute), [1]

ShardingAnalyzer (class in eformer.escale.helpers.base) ShardingRule (class in eformer.escale.helpers.base) shm_size (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] should_scale_up_multislice() (eformer.executor.ray.pool_manager.SlicePoolManager method) should_write_shared_checkpoint_files() (in module eformer.serialization.fsspec_utils) sign_radius (eformer.optimizers._config.ScionConfig attribute), [1] simulate (eformer.ops.quantization._config.QuantizationConfig attribute), [1] single_dequantize_nf4() (in module eformer.ops.quantization.quantization_functions) single_quantize_and_pack_nf4() (in module eformer.ops.quantization.quantization_functions) size (eformer.common_types.Mesh attribute) skew() (in module eformer.optimizers._tx.white_kron) SkewOptimizer (class in eformer.optimizers._builders) slice_id (eformer.executor.ray.types.MultisliceInfo attribute), [1] slice_name (eformer.executor.ray.types.HostInfo attribute), [1] (eformer.executor.ray.types.SliceInfo attribute), [1] SliceInfo (class in eformer.executor.ray.types) SlicePoolManager (class in eformer.executor.ray.pool_manager) SnitchRecipient (class in eformer.executor.ray.types) SoapConfig (class in eformer.optimizers._config) specs_to_name_sharding() (in module eformer.pytree._tree_util) spectral_radius (eformer.optimizers._config.ScionConfig attribute), [1] split() (in module eformer.pytree._tree_util) sqrt_1bit_x() (in module eformer.ops.quantization.implicit_array_1bit) sqrt_8bit_x() (in module eformer.ops.quantization.implicit_array_8bit) sqrt_nf4_x() (in module eformer.ops.quantization.implicit_array_nf4) SRowWise (class in eformer.common_types) stack (eformer.optimizers._tx.white_kron.LeafState attribute), [1] StageLocalApplyFn (in module eformer.optimizers._stage_local) StageLocalGradientTransformation (class in eformer.optimizers._stage_local) StageLocalOptimizerMetadata (class in eformer.optimizers._stage_local) start_head_node() (eformer.executor.patch_tpus_ray.RayManager method) start_time (eformer.loggings.ProgressLogger attribute) start_worker_node() (eformer.executor.patch_tpus_ray.RayManager method) start_workers (eformer.executor.cluster_util.RayClusterConfig attribute), [1] stat() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) state (eformer.executor.ray.types.JobInfo attribute), [1] StateValidationResult (class in eformer.pytree._tree_util) ste() (in module eformer.jaximus._imus) stem() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) step_policies (eformer.serialization.checkpointer.Checkpointer attribute) steps (eformer.optimizers._config.SchedulerConfig attribute), [1] stop_cluster() (eformer.executor.patch_tpus_ray.RayManager method) stop_node() (eformer.executor.patch_tpus_ray.RayManager method) straight_through() (in module eformer.ops.quantization._config) string_key_to_int() (in module eformer.pytree._tree_util) string_to_bool() (in module eformer.aparser._aparser) STRING_TO_DTYPE_MAP (in module eformer.mpric.dtypes.precision_types) suffix (eformer.paths.GCSPath property) (eformer.paths.LocalPath property) suffixes() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method)

T

target_merged_dim_size (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] tb (eformer.executor.ray.types.ExceptionInfo attribute), [1] tensor_parallel_axis (eformer.escale.partition.manager.PartitionAxis attribute) TERNARY (eformer.ops.quantization._config.QuantizationType attribute) test_ssh_connectivity() (eformer.executor.patch_tpus_ray.IPManager static method) timestamp (eformer.serialization.async_manager.CheckpointMetadata attribute), [1] to_bytes() (in module eformer.pytree._serialization) to_compressed_bytes() (in module eformer.pytree._serialization) to_dict() (eformer.escale.partition.manager.PartitionManager method) (eformer.optimizers._config.SerializationMixin method), [1] (eformer.optimizers._tx.white_kron.DenseState method) (eformer.optimizers._tx.white_kron.LeafState method) (eformer.pytree._tree_util.StateValidationResult method) (eformer.pytree._xtree.xTree method) (eformer.serialization.async_manager.CheckpointMetadata method) to_host() (in module eformer.serialization.utils) to_json() (eformer.escale.partition.manager.PartitionManager method) (eformer.optimizers._config.SerializationMixin method), [1] (eformer.optimizers._tx.white_kron.DenseState method) (eformer.optimizers._tx.white_kron.LeafState method) (eformer.pytree._tree_util.StateValidationResult method) to_kwargs() (eformer.executor.ray.resource_manager.RayResources method) to_ray_resources() (eformer.executor.ray.resource_manager.ComputeResourceConfig method) (eformer.executor.ray.resource_manager.CpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.GpuAcceleratorConfig method) (eformer.executor.ray.resource_manager.TpuAcceleratorConfig method) to_resource_dict() (eformer.executor.ray.resource_manager.RayResources method) to_state_dict() (eformer.pytree._xtree.xTree method) (in module eformer.pytree._serialization) tpu_version (eformer.executor.ray.executor.TpuRemoteManager attribute) (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] TpuAcceleratorConfig (class in eformer.executor.ray.resource_manager) TpuRemoteManager (class in eformer.executor.ray.executor) training_step_wrapper() (eformer.mpric.handler.precision_handler.PrecisionHandler method) transpose_1bit_operand() (in module eformer.ops.quantization.implicit_array_1bit) transpose_8bit_operand() (in module eformer.ops.quantization.implicit_array_8bit) transpose_nf4_operand() (in module eformer.ops.quantization.implicit_array_nf4) tree_abs() (in module eformer.pytree._tree_util) tree_add() (in module eformer.pytree._tree_util) tree_all() (in module eformer.pytree._tree_util) tree_any() (in module eformer.pytree._tree_util) tree_apply() (in module eformer.pytree._tree_util) tree_bytes() (in module eformer.pytree._tree_util) tree_cast() (in module eformer.pytree._tree_util) tree_clip() (in module eformer.pytree._tree_util) tree_concatenate() (in module eformer.pytree._tree_util) tree_deserialize_leaves() (in module eformer.serialization.serialization) tree_divide() (in module eformer.pytree._tree_util)

tree_dot() (in module eformer.pytree._tree_util) tree_equal() (in module eformer.pytree._tree_util) tree_exp() (in module eformer.pytree._tree_util) tree_expand_dims() (in module eformer.pytree._tree_util) tree_filter() (in module eformer.pytree._tree_util) tree_flatten_one_level_with_keys() (in module eformer.pytree._tree_util) tree_flatten_with_implicit() (in module eformer.jaximus._imus) tree_flatten_with_keys() (eformer.jaximus._imus.ImplicitArray method) tree_flatten_with_path_with_implicit() (in module eformer.jaximus._imus) tree_flatten_with_paths() (in module eformer.pytree._tree_util) tree_isfinite() (in module eformer.pytree._tree_util) tree_isinf() (in module eformer.pytree._tree_util) tree_isnan() (in module eformer.pytree._tree_util) tree_leaves_with_implicit() (in module eformer.jaximus._imus) tree_leaves_with_paths() (in module eformer.pytree._tree_util) tree_log() (in module eformer.pytree._tree_util) tree_map_with_implicit() (in module eformer.jaximus._imus) tree_map_with_path() (in module eformer.pytree._tree_util) tree_map_with_path_with_implicit() (in module eformer.jaximus._imus) tree_max() (in module eformer.pytree._tree_util) tree_mean() (in module eformer.pytree._tree_util) tree_min() (in module eformer.pytree._tree_util) tree_multiply() (in module eformer.pytree._tree_util) tree_norm() (in module eformer.pytree._tree_util) tree_ones_like() (in module eformer.pytree._tree_util) tree_path_to_string() (in module eformer.pytree._tree_util) tree_random_like() (in module eformer.pytree._tree_util) tree_reciprocal() (in module eformer.pytree._tree_util) tree_reduce() (in module eformer.pytree._tree_util) tree_replace_infs() (in module eformer.pytree._tree_util) tree_replace_nans() (in module eformer.pytree._tree_util) tree_reshape() (in module eformer.pytree._tree_util) tree_round() (in module eformer.pytree._tree_util) tree_serialize_leaves() (in module eformer.serialization.serialization) tree_sign() (in module eformer.pytree._tree_util) tree_size() (in module eformer.pytree._tree_util) tree_sqrt() (in module eformer.pytree._tree_util) tree_squeeze() (in module eformer.pytree._tree_util) tree_stack() (in module eformer.pytree._tree_util) tree_structure_equal() (in module eformer.pytree._tree_util) tree_structure_with_implicit() (in module eformer.jaximus._imus) tree_subtract() (in module eformer.pytree._tree_util) tree_sum() (in module eformer.pytree._tree_util) tree_transpose() (in module eformer.pytree._tree_util) tree_unflatten() (eformer.jaximus._imus.ImplicitArray class method) tree_where() (in module eformer.pytree._tree_util) tree_zeros_like() (in module eformer.pytree._tree_util) TreeFilter (class in eformer.pytree._tree_util) type_hints (eformer.pytree._pytree.PyTreeClassInfo attribute)

U

u4toi4() (in module eformer.ops.quantization.quantization_functions) unconstrained (eformer.optimizers._config.ScionConfig attribute), [1] unflatten_dict() (in module eformer.pytree._tree_util) unflatten_mapping() (in module eformer.pytree._tree_util) UnifiedExpertColumnWise (class in eformer.common_types) UnifiedExpertRowWise (class in eformer.common_types) UninitializedAval UniversalPath (class in eformer.paths) unlink() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) unpack_pytree() (in module eformer.pytree._tree_util) unpack_weights_1bit() (in module eformer.ops.quantization.quantization_functions)

unregister() (eformer.escale.partition.manager.PartitionAxis class method) unscale() (eformer.mpric.loss_scaling.loss_scaler.DynamicLossScale method) (eformer.mpric.loss_scaling.loss_scaler.NoOpLossScale method) until (eformer.serialization.checkpointer.CheckpointInterval attribute), [1] update() (eformer.common_types.Mesh method) (eformer.loggings.ProgressLogger method) update_fn_resource_env() (eformer.executor.ray.resource_manager.RayResources static method) update_prob_flat_start (eformer.optimizers._config.KronConfig attribute), [1] update_simple() (eformer.loggings.ProgressLogger method) use_implicit() (in module eformer.jaximus._imus) use_kernel (eformer.ops.quantization._config.QuantizationConfig attribute), [1] use_tensorstore (eformer.serialization.async_manager.AsyncCheckpointManager attribute) use_tty (eformer.loggings.ProgressLogger attribute) user (eformer.executor.ray.docker_executor.DockerConfig attribute), [1]

V

valid_cols (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] valid_count (eformer.optimizers._tx.white_kron.DenseState attribute), [1] valid_rows (eformer.optimizers._tx.white_kron.DenseState attribute), [1] (eformer.optimizers._tx.white_kron.LeafState attribute), [1] validate() (eformer.optimizers._base.OptimizerBuilder method), [1] validate_partition_specs() (eformer.escale.helpers.base.ShardingAnalyzer method) validate_serializable() (in module eformer.pytree._serialization) validate_sharding_config() (in module eformer.escale.partition.auto_spec)

validate_sharding_tree() (in module eformer.serialization.sharding_utils) verbose (eformer.serialization.async_manager.AsyncCheckpointManager attribute) (eformer.serialization.base_manager.CheckpointManager attribute) verify_cluster() (eformer.executor.patch_tpus_ray.RayManager method) version (eformer.serialization.async_manager.CheckpointMetadata attribute), [1] VOCAB (in module eformer.common_types) vocab_axis (eformer.escale.partition.manager.PartitionAxis attribute) volumes (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] vrn_auto_partition_spec() (in module eformer.escale.partition.auto_spec)

W

wait_until_finished() (eformer.serialization.checkpointer.Checkpointer method) warmup_steps (eformer.optimizers._config.SchedulerConfig attribute), [1] warn_on_materialize (eformer.ops.quantization.implicit_array_1bit.Array1B attribute) (eformer.ops.quantization.implicit_array_8bit.Array8B attribute) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 attribute) (eformer.ops.quantization.implicit_array_rsr.RSROperatorBinary attribute) (eformer.ops.quantization.implicit_array_rsr.RSROperatorTernary attribute) warn_once() (eformer.loggings.LazyLogger method) warning_once() (eformer.loggings.LazyLogger method) weight (eformer.ops.quantization.implicit_array_1bit.Array1B attribute), [1] (eformer.ops.quantization.implicit_array_8bit.Array8B attribute), [1] weight_decay (eformer.optimizers._config.KronConfig attribute), [1] (eformer.optimizers._config.MarsConfig attribute), [1] (eformer.optimizers._config.MuonConfig attribute), [1] (eformer.optimizers._config.SoapConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] weight_decay_mask (eformer.optimizers._config.MuonConfig attribute), [1] (eformer.optimizers._config.WhiteKronConfig attribute), [1] (eformer.optimizers._stage_local.StageLocalOptimizerMetadata attribute), [1] weight_decay_rate (eformer.optimizers._config.AdafactorConfig attribute), [1] WhiteKronConfig (class in eformer.optimizers._config) with_environment_variables() (eformer.executor.ray.resource_manager.ComputeResourceConfig method)

with_mesh_and_axis() (eformer.ops.quantization.implicit_array_8bit.Array8B method) (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method) with_name() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) with_sharding() (eformer.ops.quantization.implicit_array_8bit.Array8B method), [1] (eformer.ops.quantization.implicit_array_nf4.ArrayNF4 method), [1] with_sharding_constraint() (in module eformer.escale.partition.constraints) with_stem() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) with_suffix() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) workdir (eformer.executor.ray.docker_executor.DockerConfig attribute), [1] worker_count (eformer.executor.ray.resource_manager.CpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.GpuAcceleratorConfig attribute), [1] (eformer.executor.ray.resource_manager.TpuAcceleratorConfig attribute), [1] write_bytes() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method) write_text() (eformer.paths.GCSPath method) (eformer.paths.LocalPath method) (eformer.paths.UniversalPath method)

X

xfrom_state_dict() (in module eformer.pytree._xtree)

xto_state_dict() (in module eformer.pytree._xtree) xTree (class in eformer.pytree._xtree)