├── .gitignore ├── .readthedocs.yaml ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── docs ├── Makefile ├── _static │ └── style.css ├── api_docs │ ├── __init__.rst │ ├── aparser.rst │ ├── aparser___init__.rst │ ├── aparser__aparser.rst │ ├── apis.rst │ ├── common_types.rst │ ├── escale.rst │ ├── escale___init__.rst │ ├── escale_helpers.rst │ ├── escale_helpers___init__.rst │ ├── escale_helpers_base.rst │ ├── escale_mesh.rst │ ├── escale_mesh___init__.rst │ ├── escale_mesh_creation.rst │ ├── escale_mesh_mesh_helpers.rst │ ├── escale_partition.rst │ ├── escale_partition___init__.rst │ ├── escale_partition_auto_spec.rst │ ├── escale_partition_constraints.rst │ ├── escale_partition_manager.rst │ ├── executor.rst │ ├── executor___init__.rst │ ├── executor_cluster_util.rst │ ├── executor_docker_tpu.rst │ ├── executor_patch_tpus_ray.rst │ ├── executor_ray.rst │ ├── executor_ray___init__.rst │ ├── executor_ray_docker_executor.rst │ ├── executor_ray_executor.rst │ ├── executor_ray_pool_manager.rst │ ├── executor_ray_resource_manager.rst │ ├── executor_ray_types.rst │ ├── jaximus.rst │ ├── jaximus___init__.rst │ ├── jaximus__imus.rst │ ├── loggings.rst │ ├── mpric.rst │ ├── mpric___init__.rst │ ├── mpric_dtypes.rst │ ├── mpric_dtypes___init__.rst │ ├── mpric_dtypes_precision_types.rst │ ├── mpric_handler.rst │ ├── mpric_handler___init__.rst │ ├── mpric_handler_precision_handler.rst │ ├── mpric_loss_scaling.rst │ ├── mpric_loss_scaling___init__.rst │ ├── mpric_loss_scaling_loss_scaler.rst │ ├── mpric_policy.rst │ ├── mpric_policy___init__.rst │ ├── mpric_policy_policy.rst │ ├── ops.rst │ ├── ops___init__.rst │ ├── ops_quantization.rst │ ├── ops_quantization___init__.rst │ ├── ops_quantization_implicit_array_1bit.rst │ ├── ops_quantization_implicit_array_8bit.rst │ ├── ops_quantization_implicit_array_nf4.rst │ ├── ops_quantization_implicit_array_rsr.rst │ ├── ops_quantization_quantization_functions.rst │ ├── optimizers.rst │ ├── optimizers___init__.rst │ ├── optimizers__base.rst │ ├── optimizers__builders.rst │ ├── optimizers__config.rst │ ├── optimizers__factory.rst │ ├── optimizers__tx.rst │ ├── optimizers__tx___init__.rst │ ├── optimizers__tx_mars.rst │ ├── optimizers__tx_utils.rst │ ├── optimizers__tx_white_kron.rst │ ├── paths.rst │ ├── pytree.rst │ ├── pytree___init__.rst │ ├── pytree__pytree.rst │ ├── pytree__serialization.rst │ ├── pytree__tree_util.rst │ ├── pytree__xtree.rst │ ├── serialization.rst │ ├── serialization___init__.rst │ ├── serialization_async_manager.rst │ ├── serialization_base_manager.rst │ ├── serialization_checkpointer.rst │ ├── serialization_fsspec_utils.rst │ ├── serialization_serialization.rst │ ├── serialization_sharding_utils.rst │ └── serialization_utils.rst ├── conf.py ├── contributing.rst ├── index.rst ├── make.bat ├── pytree_utils.md └── requirements.txt ├── eformer ├── __init__.py ├── aparser │ ├── __init__.py │ └── _aparser.py ├── common_types.py ├── escale │ ├── __init__.py │ ├── helpers │ │ ├── __init__.py │ │ └── base.py │ ├── mesh │ │ ├── __init__.py │ │ ├── creation.py │ │ └── mesh_helpers.py │ └── partition │ │ ├── __init__.py │ │ ├── auto_spec.py │ │ ├── constraints.py │ │ └── manager.py ├── executor │ ├── __init__.py │ ├── cluster_util.py │ ├── docker_tpu.py │ ├── patch_tpus_ray.py │ └── ray │ │ ├── __init__.py │ │ ├── docker_executor.py │ │ ├── executor.py │ │ ├── pool_manager.py │ │ ├── resource_manager.py │ │ └── types.py ├── jaximus │ ├── __init__.py │ └── _imus.py ├── loggings.py ├── mpric │ ├── __init__.py │ ├── dtypes │ │ ├── __init__.py │ │ └── precision_types.py │ ├── handler │ │ ├── __init__.py │ │ └── precision_handler.py │ ├── loss_scaling │ │ ├── __init__.py │ │ └── loss_scaler.py │ └── policy │ │ ├── __init__.py │ │ └── policy.py ├── ops │ ├── __init__.py │ └── quantization │ │ ├── __init__.py │ │ ├── _config.py │ │ ├── implicit_array_1bit.py │ │ ├── implicit_array_8bit.py │ │ ├── implicit_array_nf4.py │ │ ├── implicit_array_rsr.py │ │ └── quantization_functions.py ├── optimizers │ ├── __init__.py │ ├── _base.py │ ├── _builders.py │ ├── _config.py │ ├── _factory.py │ └── _tx │ │ ├── __init__.py │ │ ├── mars.py │ │ ├── utils.py │ │ └── white_kron.py ├── paths.py ├── py.typed ├── pytree │ ├── __init__.py │ ├── _pytree.py │ ├── _serialization.py │ ├── _tree_util.py │ └── _xtree.py └── serialization │ ├── __init__.py │ ├── async_manager.py │ ├── base_manager.py │ ├── checkpointer.py │ ├── fsspec_utils.py │ ├── serialization.py │ ├── sharding_utils.py │ └── utils.py ├── examples └── quantization_training.py ├── pre_commit.py ├── pyproject.toml ├── tests ├── test_autoptree.py ├── test_optimizers.py ├── test_quantization.py └── test_tree_utils.py └── uv.lock /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/.gitignore -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/.readthedocs.yaml -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/README.md -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/_static/style.css: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/api_docs/__init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/__init__.rst -------------------------------------------------------------------------------- /docs/api_docs/aparser.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/aparser.rst -------------------------------------------------------------------------------- /docs/api_docs/aparser___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/aparser___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/aparser__aparser.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/aparser__aparser.rst -------------------------------------------------------------------------------- /docs/api_docs/apis.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/apis.rst -------------------------------------------------------------------------------- /docs/api_docs/common_types.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/common_types.rst -------------------------------------------------------------------------------- /docs/api_docs/escale.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale.rst -------------------------------------------------------------------------------- /docs/api_docs/escale___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_helpers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_helpers.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_helpers___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_helpers___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_helpers_base.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_helpers_base.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_mesh.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_mesh.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_mesh___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_mesh___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_mesh_creation.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_mesh_creation.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_mesh_mesh_helpers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_mesh_mesh_helpers.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_partition.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_partition.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_partition___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_partition___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_partition_auto_spec.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_partition_auto_spec.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_partition_constraints.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_partition_constraints.rst -------------------------------------------------------------------------------- /docs/api_docs/escale_partition_manager.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/escale_partition_manager.rst -------------------------------------------------------------------------------- /docs/api_docs/executor.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor.rst -------------------------------------------------------------------------------- /docs/api_docs/executor___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_cluster_util.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_cluster_util.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_docker_tpu.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_docker_tpu.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_patch_tpus_ray.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_patch_tpus_ray.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray_docker_executor.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray_docker_executor.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray_executor.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray_executor.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray_pool_manager.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray_pool_manager.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray_resource_manager.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray_resource_manager.rst -------------------------------------------------------------------------------- /docs/api_docs/executor_ray_types.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/executor_ray_types.rst -------------------------------------------------------------------------------- /docs/api_docs/jaximus.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/jaximus.rst -------------------------------------------------------------------------------- /docs/api_docs/jaximus___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/jaximus___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/jaximus__imus.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/jaximus__imus.rst -------------------------------------------------------------------------------- /docs/api_docs/loggings.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/loggings.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_dtypes.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_dtypes.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_dtypes___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_dtypes___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_dtypes_precision_types.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_dtypes_precision_types.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_handler.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_handler.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_handler___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_handler___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_handler_precision_handler.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_handler_precision_handler.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_loss_scaling.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_loss_scaling.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_loss_scaling___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_loss_scaling___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_loss_scaling_loss_scaler.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_loss_scaling_loss_scaler.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_policy.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_policy.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_policy___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_policy___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/mpric_policy_policy.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/mpric_policy_policy.rst -------------------------------------------------------------------------------- /docs/api_docs/ops.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops.rst -------------------------------------------------------------------------------- /docs/api_docs/ops___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization_implicit_array_1bit.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization_implicit_array_1bit.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization_implicit_array_8bit.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization_implicit_array_8bit.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization_implicit_array_nf4.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization_implicit_array_nf4.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization_implicit_array_rsr.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization_implicit_array_rsr.rst -------------------------------------------------------------------------------- /docs/api_docs/ops_quantization_quantization_functions.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/ops_quantization_quantization_functions.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__base.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__base.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__builders.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__builders.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__config.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__config.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__factory.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__factory.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__tx.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__tx.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__tx___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__tx___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__tx_mars.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__tx_mars.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__tx_utils.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__tx_utils.rst -------------------------------------------------------------------------------- /docs/api_docs/optimizers__tx_white_kron.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/optimizers__tx_white_kron.rst -------------------------------------------------------------------------------- /docs/api_docs/paths.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/paths.rst -------------------------------------------------------------------------------- /docs/api_docs/pytree.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/pytree.rst -------------------------------------------------------------------------------- /docs/api_docs/pytree___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/pytree___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/pytree__pytree.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/pytree__pytree.rst -------------------------------------------------------------------------------- /docs/api_docs/pytree__serialization.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/pytree__serialization.rst -------------------------------------------------------------------------------- /docs/api_docs/pytree__tree_util.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/pytree__tree_util.rst -------------------------------------------------------------------------------- /docs/api_docs/pytree__xtree.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/pytree__xtree.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization___init__.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization___init__.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_async_manager.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_async_manager.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_base_manager.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_base_manager.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_checkpointer.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_checkpointer.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_fsspec_utils.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_fsspec_utils.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_serialization.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_serialization.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_sharding_utils.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_sharding_utils.rst -------------------------------------------------------------------------------- /docs/api_docs/serialization_utils.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/api_docs/serialization_utils.rst -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/conf.py -------------------------------------------------------------------------------- /docs/contributing.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/contributing.rst -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/index.rst -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/make.bat -------------------------------------------------------------------------------- /docs/pytree_utils.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/pytree_utils.md -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/docs/requirements.txt -------------------------------------------------------------------------------- /eformer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/__init__.py -------------------------------------------------------------------------------- /eformer/aparser/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/aparser/__init__.py -------------------------------------------------------------------------------- /eformer/aparser/_aparser.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/aparser/_aparser.py -------------------------------------------------------------------------------- /eformer/common_types.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/common_types.py -------------------------------------------------------------------------------- /eformer/escale/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/__init__.py -------------------------------------------------------------------------------- /eformer/escale/helpers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/helpers/__init__.py -------------------------------------------------------------------------------- /eformer/escale/helpers/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/helpers/base.py -------------------------------------------------------------------------------- /eformer/escale/mesh/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/mesh/__init__.py -------------------------------------------------------------------------------- /eformer/escale/mesh/creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/mesh/creation.py -------------------------------------------------------------------------------- /eformer/escale/mesh/mesh_helpers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/mesh/mesh_helpers.py -------------------------------------------------------------------------------- /eformer/escale/partition/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/partition/__init__.py -------------------------------------------------------------------------------- /eformer/escale/partition/auto_spec.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/partition/auto_spec.py -------------------------------------------------------------------------------- /eformer/escale/partition/constraints.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/partition/constraints.py -------------------------------------------------------------------------------- /eformer/escale/partition/manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/escale/partition/manager.py -------------------------------------------------------------------------------- /eformer/executor/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/__init__.py -------------------------------------------------------------------------------- /eformer/executor/cluster_util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/cluster_util.py -------------------------------------------------------------------------------- /eformer/executor/docker_tpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/docker_tpu.py -------------------------------------------------------------------------------- /eformer/executor/patch_tpus_ray.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/patch_tpus_ray.py -------------------------------------------------------------------------------- /eformer/executor/ray/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/ray/__init__.py -------------------------------------------------------------------------------- /eformer/executor/ray/docker_executor.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/ray/docker_executor.py -------------------------------------------------------------------------------- /eformer/executor/ray/executor.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/ray/executor.py -------------------------------------------------------------------------------- /eformer/executor/ray/pool_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/ray/pool_manager.py -------------------------------------------------------------------------------- /eformer/executor/ray/resource_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/ray/resource_manager.py -------------------------------------------------------------------------------- /eformer/executor/ray/types.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/executor/ray/types.py -------------------------------------------------------------------------------- /eformer/jaximus/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/jaximus/__init__.py -------------------------------------------------------------------------------- /eformer/jaximus/_imus.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/jaximus/_imus.py -------------------------------------------------------------------------------- /eformer/loggings.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/loggings.py -------------------------------------------------------------------------------- /eformer/mpric/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/__init__.py -------------------------------------------------------------------------------- /eformer/mpric/dtypes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/dtypes/__init__.py -------------------------------------------------------------------------------- /eformer/mpric/dtypes/precision_types.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/dtypes/precision_types.py -------------------------------------------------------------------------------- /eformer/mpric/handler/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/handler/__init__.py -------------------------------------------------------------------------------- /eformer/mpric/handler/precision_handler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/handler/precision_handler.py -------------------------------------------------------------------------------- /eformer/mpric/loss_scaling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/loss_scaling/__init__.py -------------------------------------------------------------------------------- /eformer/mpric/loss_scaling/loss_scaler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/loss_scaling/loss_scaler.py -------------------------------------------------------------------------------- /eformer/mpric/policy/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/policy/__init__.py -------------------------------------------------------------------------------- /eformer/mpric/policy/policy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/mpric/policy/policy.py -------------------------------------------------------------------------------- /eformer/ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/__init__.py -------------------------------------------------------------------------------- /eformer/ops/quantization/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/__init__.py -------------------------------------------------------------------------------- /eformer/ops/quantization/_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/_config.py -------------------------------------------------------------------------------- /eformer/ops/quantization/implicit_array_1bit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/implicit_array_1bit.py -------------------------------------------------------------------------------- /eformer/ops/quantization/implicit_array_8bit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/implicit_array_8bit.py -------------------------------------------------------------------------------- /eformer/ops/quantization/implicit_array_nf4.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/implicit_array_nf4.py -------------------------------------------------------------------------------- /eformer/ops/quantization/implicit_array_rsr.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/implicit_array_rsr.py -------------------------------------------------------------------------------- /eformer/ops/quantization/quantization_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/ops/quantization/quantization_functions.py -------------------------------------------------------------------------------- /eformer/optimizers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/__init__.py -------------------------------------------------------------------------------- /eformer/optimizers/_base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_base.py -------------------------------------------------------------------------------- /eformer/optimizers/_builders.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_builders.py -------------------------------------------------------------------------------- /eformer/optimizers/_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_config.py -------------------------------------------------------------------------------- /eformer/optimizers/_factory.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_factory.py -------------------------------------------------------------------------------- /eformer/optimizers/_tx/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_tx/__init__.py -------------------------------------------------------------------------------- /eformer/optimizers/_tx/mars.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_tx/mars.py -------------------------------------------------------------------------------- /eformer/optimizers/_tx/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_tx/utils.py -------------------------------------------------------------------------------- /eformer/optimizers/_tx/white_kron.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/optimizers/_tx/white_kron.py -------------------------------------------------------------------------------- /eformer/paths.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/paths.py -------------------------------------------------------------------------------- /eformer/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /eformer/pytree/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/pytree/__init__.py -------------------------------------------------------------------------------- /eformer/pytree/_pytree.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/pytree/_pytree.py -------------------------------------------------------------------------------- /eformer/pytree/_serialization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/pytree/_serialization.py -------------------------------------------------------------------------------- /eformer/pytree/_tree_util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/pytree/_tree_util.py -------------------------------------------------------------------------------- /eformer/pytree/_xtree.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/pytree/_xtree.py -------------------------------------------------------------------------------- /eformer/serialization/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/__init__.py -------------------------------------------------------------------------------- /eformer/serialization/async_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/async_manager.py -------------------------------------------------------------------------------- /eformer/serialization/base_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/base_manager.py -------------------------------------------------------------------------------- /eformer/serialization/checkpointer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/checkpointer.py -------------------------------------------------------------------------------- /eformer/serialization/fsspec_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/fsspec_utils.py -------------------------------------------------------------------------------- /eformer/serialization/serialization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/serialization.py -------------------------------------------------------------------------------- /eformer/serialization/sharding_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/sharding_utils.py -------------------------------------------------------------------------------- /eformer/serialization/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/eformer/serialization/utils.py -------------------------------------------------------------------------------- /examples/quantization_training.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/examples/quantization_training.py -------------------------------------------------------------------------------- /pre_commit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/pre_commit.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/pyproject.toml -------------------------------------------------------------------------------- /tests/test_autoptree.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/tests/test_autoptree.py -------------------------------------------------------------------------------- /tests/test_optimizers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/tests/test_optimizers.py -------------------------------------------------------------------------------- /tests/test_quantization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/tests/test_quantization.py -------------------------------------------------------------------------------- /tests/test_tree_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/tests/test_tree_utils.py -------------------------------------------------------------------------------- /uv.lock: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/erfanzar/eformer/HEAD/uv.lock --------------------------------------------------------------------------------