Views
No views yet
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/gemma-4-E4B"
5
6dtype = jnp.bfloat16 # try jnp.float16 on many GPUs
7
8model = ed.AutoEasyDeLModelForImageTextToText.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)from_torch=True to from_pretrained(...).dp: data parallel (replicated parameters, different batch shards)fsdp: parameter sharding (memory saver; often the biggest axis)ep: expert parallel (MoE; keep 1 for non-MoE models)tp: tensor parallel (splits large matmuls)sp: sequence parallel (splits sequence dimension)sharding_axis_names=("dp","fsdp","ep","tp","sp") and choose sharding_axis_dims so that their product equals your device count.
You can use -1 in sharding_axis_dims to let EasyDeL infer the remaining dimension.1# 8 devices, pure FSDP
2sharding_axis_dims = (1, 8, 1, 1, 1)
3
4# 8 devices, 2-way DP x 4-way FSDP
5sharding_axis_dims = (2, 4, 1, 1, 1)
6
7# 8 devices, 4-way FSDP x 2-way TP
8sharding_axis_dims = (1, 4, 1, 2, 1)eLargeModel (ELM)eLargeModel is a higher-level interface that wires together loading, sharding, training, and eSurge inference from a single config.1from easydel import eLargeModel
2
3repo_id = "EasyDeL/gemma-4-E4B"
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()1model:
2 name_or_path: "EasyDeL/gemma-4-E4B"
3
4loader:
5 dtype: bf16
6 param_dtype: bf16
7
8sharding:
9 axis_dims: [1, -1, 1, 1, 1]
10 auto_shard_model: trueAttentionMechanisms.*dtype, param_dtype, and precisionpip install easydel1@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}