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

Training Run: olmo3-7b-think-posttraining8k

Run Summary

Key	Value
Model	olmo3-7b-think-posttraining8k
Architecture	olmo3
Task	CausalLM
Attention	auto (AttentionMechanisms.AUTO)
Platform	TPU
Devices (total/local)	32 / 4
Dtype / Param dtype	<class 'jax.numpy.bfloat16'> / <class 'jax.numpy.bfloat16'>
EasyDeL	v0.3.0

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.

Example sharding configs

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

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

Key	Value
Learning rate	2.5e-05 -> 0.0
Optimizer	adamw
Scheduler	linear
Warmup steps	5
Weight decay	0.0
Loss config	`LossConfig(
ignore_index : -100
label_smoothing : 0.0
z_loss : 0.0
loss_normalizing_factor : NUM_REAL_TARGET_TOKENS
num_labels : None
problem_type : None
divide_weight_sum : False
shift_tokens : True
break_on_nan : True
reduction : None
num_classification_labels : None
classification_problem_type : None
chunk_vocab_size : None
chunk_token_size : None
chunk_block_size : 4096
compute_dtype : 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

Partition rules

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

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.