# 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*