# 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**: 5
- **Batch Size**: 16
- **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*