This checkpoint is intended to be loaded with EasyDeL on JAX (CPU/GPU/TPU). It supports sharded loading with auto_shard_model=True and configurable precision via dtype, param_dtype, and precision.
1import easydel as ed
2from jax import numpy as jnp, lax
3
4repo_id = "EasyDeL/Llama-3.1-8B"
5
6dtype = jnp.bfloat16 # try jnp.float16 on many GPUs
7
8model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
9 repo_id,
10 dtype=dtype,
11 param_dtype=dtype,
12 precision=lax.Precision("fastest"),
13 sharding_axis_names=("dp", "fsdp", "ep", "tp", "sp"),
14 sharding_axis_dims=(1, -1, 1, 1, 1),
15 config_kwargs=ed.EasyDeLBaseConfigDict(
16 attn_dtype=dtype,
17 attn_mechanism=ed.AttentionMechanisms.VANILLA,
18 fsdp_is_ep_bound=True,
19 sp_is_ep_bound=True,
20 moe_method=ed.MoEMethods.FUSED_MOE,
21 ),
22 auto_shard_model=True,
23 partition_axis=ed.PartitionAxis(),
24)
EasyDeL can scale to multiple devices by creating a logical device mesh. Most EasyDeL loaders use a 5D mesh:
1from easydel import eLargeModel
2
3repo_id = "EasyDeL/Llama-3.1-8B"
4
5elm = eLargeModel.from_pretrained(repo_id) # task is auto-detected
6elm.set_dtype("bf16")
7elm.set_sharding(axis_names=("dp", "fsdp", "ep", "tp", "sp"), axis_dims=(1, -1, 1, 1, 1))
8
9model = elm.build_model()
10# Optional: build an inference engine
11# engine = elm.build_esurge()
EasyDeL is released under the Apache-2.0 license. The license for this model's weights may differ; please consult the original repository.
1@misc{Zare Chavoshi_2023,
2 title={EasyDeL: An open-source library for enhancing and streamlining the training process of machine learning models},
3 url={https://github.com/erfanzar/EasyDeL},
4 author={Zare Chavoshi, Erfan},
5 year={2023}
6}
