| # 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**: qwen2 | |
| - **Platform**: TPU | |
| - **Number of Devices**: 16 | |
| ### Training Parameters | |
| - **Learning Rate**: 5e-05 → 5e-06 | |
| - **Optimizer**: adamw | |
| - **Scheduler**: cosine | |
| - **Warmup Steps**: 160 | |
| - **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**: 2 | |
| - **Batch Size**: 8 | |
| - **Sequence Length**: 4096 | |
| - **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('tp', ('fsdp', 'sp'))), | |
| ( 'self_attn/(q_proj|k_proj|v_proj)/kernel', | |
| PartitionSpec(('fsdp', 'sp'), 'tp')), | |
| ('self_attn/o_proj/kernel', PartitionSpec('tp', ('fsdp', 'sp'))), | |
| ('mlp/gate_proj/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), | |
| ('mlp/down_proj/kernel', PartitionSpec('tp', ('fsdp', 'sp'))), | |
| ('mlp/up_proj/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), | |
| ('input_layernorm/kernel', PartitionSpec(None,)), | |
| ('post_attention_layernorm/kernel', PartitionSpec(None,)), | |
| ('model/norm/kernel', PartitionSpec(None,)), | |
| ('lm_head/kernel', PartitionSpec(('fsdp', 'sp'), 'tp')), | |
| ('.*', PartitionSpec(None,))) | |
| ``` | |
| --- | |
| *Generated with EasyDeL v0.1.2* | |