Olmo-3-7B-PostTraining8K-SFT-step-80

<p align="center"> <a href="https://github.com/erfanzar/EasyDeL"> <img src="https://raw.githubusercontent.com/erfanzar/easydel/main/images/easydel-logo-with-text.png" height="80" alt="EasyDeL" /> </a> </p> <h1 align="center">Training Run: olmo3-7b-think-posttraining8k</h1> <p align="center"> <a href="https://github.com/erfanzar/EasyDeL"> <img src="https://img.shields.io/static/v1?label=EasyDeL&message=v0.3.0&color=blue&style=flat-square" alt="EasyDeL Version" /> </a> <img src="https://img.shields.io/static/v1?label=Arch&message=olmo3&color=0A66C2&style=flat-square" alt="Model Architecture" /> <img src="https://img.shields.io/static/v1?label=Task&message=CausalLM&color=green&style=flat-square" alt="Task" /> <img src="https://img.shields.io/static/v1?label=Attention&message=auto&color=8A2BE2&style=flat-square" alt="Attention Mechanism" /> </p>

Run Summary

Load This Checkpoint

import easydel as ed
from jax import numpy as jnp, lax

repo_id = "user/model-id"  # TODO: set this to your local output directory or HF repo ID

dtype = jnp.bfloat16
param_dtype = jnp.bfloat16

model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
    repo_id,
    config_kwargs=ed.EasyDeLBaseConfigDict(
        attn_dtype=dtype,
        attn_mechanism=ed.AttentionMechanisms.AUTO,
    ),
    dtype=dtype,
    param_dtype=param_dtype,
    precision=lax.Precision("fastest"),
    auto_shard_model=True,
)

Sharding Notes

EasyDeL typically shards across a 5D logical mesh: ("dp","fsdp","ep","tp","sp").

• The product of sharding_axis_dims should match your device count; you can use -1 to infer the remaining dimension.
• fsdp is commonly the largest axis to reduce memory usage.
• For non-MoE models keep ep=1.

<details> <summary>Example sharding configs</summary>

8 devices: (1, 8, 1, 1, 1)  # pure FSDP
8 devices: (2, 4, 1, 1, 1)  # 2-way DP x 4-way FSDP
8 devices: (1, 4, 1, 2, 1)  # 4-way FSDP x 2-way TP

</details>

Using via eLargeModel (ELM)

from easydel import eLargeModel

elm = eLargeModel.from_pretrained(repo_id)
elm.set_dtype("bf16")
elm.set_sharding(axis_names=("dp", "fsdp", "ep", "tp", "sp"), axis_dims=(1, -1, 1, 1, 1))

model = elm.build_model()
# engine = elm.build_esurge()

Hyperparameters

<details> <summary>Partition rules</summary>

( ('^(?:.*/)?lm_head/kernel(?:/.*)?$', P(('fsdp', 'sp'), 'tp')),
  ('^(?:.*/)?model/embed_tokens/embedding(?:/.*)?$', P(('fsdp', 'sp'), 'tp')),
  ( '^(?:.*/)?model/layers/\\d+/mlp/down_proj/kernel(?:/.*)?$',
    P('tp', ('fsdp', 'sp'))),
  ( '^(?:.*/)?model/layers/\\d+/mlp/gate_proj/kernel(?:/.*)?$',
    P(('fsdp', 'sp'), 'tp')),
  ( '^(?:.*/)?model/layers/\\d+/mlp/up_proj/kernel(?:/.*)?$',
    P(('fsdp', 'sp'), 'tp')),
  ( '^(?:.*/)?model/layers/\\d+/post_attention_layernorm/kernel(?:/.*)?$',
    P(None,)),
  ( '^(?:.*/)?model/layers/\\d+/post_feedforward_layernorm/kernel(?:/.*)?$',
    P(None,)),
  ('^(?:.*/)?model/layers/\\d+/self_attn/k_norm/kernel(?:/.*)?$', P(None,)),
  ( '^(?:.*/)?model/layers/\\d+/self_attn/k_proj/kernel(?:/.*)?$',
    P(('fsdp', 'sp'), 'tp')),
  ( '^(?:.*/)?model/layers/\\d+/self_attn/o_proj/kernel(?:/.*)?$',
    P('tp', ('fsdp', 'sp'))),
  ('^(?:.*/)?model/layers/\\d+/self_attn/q_norm/kernel(?:/.*)?$', P(None,)),
  ( '^(?:.*/)?model/layers/\\d+/self_attn/q_proj/kernel(?:/.*)?$',
    P(('fsdp', 'sp'), 'tp')),
  ( '^(?:.*/)?model/layers/\\d+/self_attn/v_proj/kernel(?:/.*)?$',
    P(('fsdp', 'sp'), 'tp')),
  ('^(?:.*/)?model/norm/kernel(?:/.*)?$', P(None,)),
  ('.*', P()))

</details>

Citation

@misc{Zare Chavoshi_2023,
    title={EasyDeL: An open-source library for enhancing and streamlining the training process of machine learning models},
    url={https://github.com/erfanzar/EasyDeL},
    author={Zare Chavoshi, Erfan},
    year={2023}
}

---

Generated by EasyDeL v0.3.0.