|
|
| # BaseTrainer |
|
|
| ## 🚀 Trained With [EasyDeL](https://github.com/erfanzar/EasyDeL) |
|
|
| EasyDeL is an open-source framework designed to enhance and streamline the training process of machine learning |
| models. With a primary focus on Jax, EasyDeL aims to provide convenient and effective solutions for |
| training Flax/Jax models on TPU/GPU, for both serving and training purposes. |
|
|
| ## 📦 Installation & Usage |
| |
| ```python |
| from easydel import AutoEasyDeLModelForCausalLM |
| from jax import numpy as jnp, lax |
| |
| model = AutoEasyDeLModelForCausalLM.from_pretrained( |
| f"REPO_ID/BaseTrainer", |
| dtype=..., |
| param_dtype=..., |
| precision=lax.Precision("fastest"), |
| auto_shard_model=True, |
| ) |
| ``` |
|
|
| ## 🔧 Training Configuration |
|
|
| ### Model Details |
| - **Architecture**: gemma3_text |
| - **Platform**: TPU |
| - **Number of Devices**: 16 |
| |
| ### Training Parameters |
| - **Learning Rate**: 4e-05 → 4e-06 |
| - **Optimizer**: adamw |
| - **Scheduler**: cosine |
| - **Warmup Steps**: 50 |
| - **Weight Decay**: 0.02 |
| - **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 |
| ) |
|
|
| ### Training Setup |
| - **Epochs**: 3 |
| - **Batch Size**: 8 |
| - **Sequence Length**: 8192 |
| - **Dtype**: <class 'jax.numpy.bfloat16'> |
| - **Params Dtype**: <class 'jax.numpy.bfloat16'> |
|
|
| ### Advanced Configuration |
| - **Gradient Checkpointing**: |
| - **Gradient Accumulation Steps**: 1 |
| - **Max Training Steps**: None |
| - **Max Evaluation Steps**: None |
| - **Training Duration**: 7H |
|
|
| ### Sharding Configuration |
| ```python |
| # Partition Rules |
| ( ('model/embed_tokens/embedding', PartitionSpec(('fsdp', 'sp'), 'tp')), |
| ('self_attn/q_proj/kernel', PartitionSpec('tp', ('fsdp', 'sp'))), |
| ('self_attn/k_proj/kernel', PartitionSpec('tp', ('fsdp', 'sp'))), |
| ('self_attn/v_proj/kernel', PartitionSpec('tp', ('fsdp', 'sp'))), |
| ('self_attn/o_proj/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), |
| ('mlp/gate_proj/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), |
| ('mlp/up_proj/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), |
| ('mlp/down_proj/kernel', PartitionSpec('tp', ('fsdp', 'sp'))), |
| ('input_layernorm/kernel', PartitionSpec(None,)), |
| ('post_attention_layernorm/kernel', PartitionSpec(None,)), |
| ('pre_feedforward_layernorm/kernel', PartitionSpec(None,)), |
| ('post_feedforward_layernorm/kernel', PartitionSpec(None,)), |
| ('model/norm/kernel', PartitionSpec(None,)), |
| ('lm_head/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), |
| ('.*', PartitionSpec(None,))) |
| ``` |
|
|
| --- |
| *Generated with EasyDeL v0.1.3* |
|
|
