CoolFace
Modelpublic

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

sourceHugging Faceupdated 20d agoView on Hugging Face
0likes46downloads
Model Card

<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

KeyValue
Modelolmo3-7b-think-posttraining8k
Architectureolmo3
TaskCausalLM
Attentionauto (AttentionMechanisms.AUTO)
PlatformTPU
Devices (total/local)32 / 4
Dtype / Param dtype<class 'jax.numpy.bfloat16'> / <class 'jax.numpy.bfloat16'>
EasyDeLv0.3.0

Load This Checkpoint

python
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>

text
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)

python
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

KeyValue
Learning rate2.5e-05 -> 0.0
Optimizeradamw
Schedulerlinear
Warmup steps5
Weight decay0.0
Loss config`LossConfig(

ignoreindex : -100 labelsmoothing : 0.0 zloss : 0.0 lossnormalizingfactor : NUMREALTARGETTOKENS numlabels : None problemtype : None divideweightsum : False shifttokens : True breakonnan : True reduction : None numclassificationlabels : None classificationproblemtype : None chunkvocabsize : None chunktokensize : None chunkblocksize : 4096 computedtype : fp32 ) | | Epochs | 1 | | Total batch size | 8 | | Max length | 32768 | | Grad accumulation | 1 | | Gradient checkpointing | mlp_notsaveable | | Max training steps | 118 | | Max eval steps | Not Set | | Training time limit | Not Set` |

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

text
( ('^(?:.*/)?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

bibtex
@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.