# EasyDeL 🔮
EasyDeL, an open-source library, is specifically designed to enhance and streamline the training process of machine
learning models. It focuses primarily on Jax/Flax and aims to provide convenient and effective solutions for training
Flax/Jax Models on TPU/GPU for both Serving and Training purposes. Additionally, EasyDeL will support mojo and will be
rewritten for mojo as well.
Some of the key features provided by EasyDeL include:
- DPOTrainer, SFTTrainer, and VideoCLM Trainers
- Serving and API Engines for Using and serving LLMs in JAX as efficiently as possible.
- Support Quantization Methods for all the Models.
- Support for 8, 6, and 4 BIT Operation, for inference and training in JAX
- A wide range of models in Jax is supported which have never been implemented before such as Falcon, Qwen2, Phi2,
Mixtral, Qwen2Moe, Cohere,and MPT ...
- Integration of flashAttention in JAX for GPUs and TPUs
- Automatic serving of LLMs with mid and high-level APIs in both JAX and PyTorch
- LLM Trainer and fine-tuner in JAX
- Video CLM Trainer and Fine-tuner for Models such Falcon, Qwen2, Phi2, MPT, Mixtral, Grok-1, and Qwen2Moe ...
- RLHF (Reinforcement Learning from Human Feedback) in Jax (Beta Stage)
- Various other features to enhance the training process and optimize performance.
- LoRA: Low-Rank Adaptation of Large Language Models
- RingAttention, Flash Attention, BlockWise FFN, and Efficient Attention are supported for more than 90 % of models
([FJFormer](https://github.com/erfanzar/FJFormer) Backbone).
- Serving and API Engines for Using and serving LLMs in JAX as efficient as possible.
- Automatic Converting Models from JAX-EasyDeL to PyTorch-HF and reverse
> **News**
>
> Phi3 Model bugs are fixed, Arctic Model is added.
>
> Phi3 Model is present Now Apr 24 2024
>
> Drbx Model is present
> Now [Apr 23 2024](https://github.com/erfanzar/EasyDeL/commit/4c1c5af099dad9334a82808eb04b03e7e567ddb7)
>
> New Attention Types are added `sharded_vanilla`, `wise_ring`, `sharded_vanilla` is same as `vanilla` but
> uses shard_map which will make it a little faster and more efficent.
>
> `local_ring` is Added which is Ring Attention but for TPU/GPU/CPU(s) and support attention bias instead of attention
> mask, `normal` attention is renamed to `vanilla`.
>
> `load_in_8bit` is now available for all the models, and requires to upgrade _fjformer to 0.0.50_
>
> Sharing Key and Value Cache for Large Sequence Length across devices are now Fixed (Attention Models).
>
> Cohere Model added [Apr 14 2024](https://github.com/erfanzar/EasyDeL/commit/06acb4a1afd7b67982a88b50840b90e73b1c9850)
## Documentation 💫
> [!IMPORTANT]
> Documents and Examples are ready at [Here](https://erfanzar.github.io/EasyDeL)
> Please have that in mind that EasyDel is in the loop of fast-development
> so we might have API changes.
### Hands on Code Kaggle Examples
1. [script](https://www.kaggle.com/citifer/easydel-causal-language-model-trainer-example) for mindset of using EasyDeL
CausalLanguageModelTrainer on kaggle, but you can do much more.
2. [script](https://www.kaggle.com/code/citifer/easydel-serve-example-mixtral) for using and serving LLMs with EasyDeL
JAXServer API (Mixtral Example).
3. [script](https://www.kaggle.com/code/citifer/easydel-sfttrainer-example) SuperVised Finetuning with EasyDeL.
## Serving
you can read docs or examples to see how `JAXServer` works but let me show you how you can simply host and serve any
model that supported by `EasyDeL` fo this example ill just use `gemma-7-it` by google, but you can use any model as you
wish.
```shell
python -m examples.jax_serve_example \
--prompter_type="gemma" \
--share_gradio=True \
--sharding_axis_dims=1,1,1,-1 \
--attn_mechanism="sharded_vanilla" \
--scan_ring_attention=True \
--max_sequence_length=8192 \
--max_new_tokens_ratio=25 \
--max_compile_tokens=256 \
--block_k=128 \
--block_q=128 \
--pretrained_model_name_or_path="google/gemma-7b-it" \
--dtype="bf16"
```
> [!NOTE]
> you can use `EasyServe` which is a Serve API Engine for production purpose sice that's more stable provide versioned
> API and efficient.
## Supervised Fine-Tuning with EasyDeL
EasyDeL supports both DPO and SFT Trainers, so dealing with LLMs in jax is a lot easier right now
let have an example of using Supervised Fine-Tuner in JAX with EasyDeL
```python
from EasyDel import (
TrainArguments,
AutoEasyDelModelForCausalLM,
EasyDelOptimizers,
EasyDelSchedulers,
EasyDelGradientCheckPointers,
SFTTrainer,
conversations_formatting_function # i have added this one for newcomers so if they
# don't know what's going on they can use this pre created prompter
)
from datasets import load_dataset
import flax
from jax import numpy as jnp
from transformers import AutoTokenizer
huggingface_repo_id_or_path = "mistralai/Mistral-7B-Instruct-v0.2"
model, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )
max_length = 4096
tokenizer = AutoTokenizer.from_pretrained(
huggingface_repo_id_or_path,
trust_remote_code=True
)
tokenizer.pad_token = tokenizer.eos_token
configs_to_initialize_model_class = {
"config": model.config,
"dtype": jnp.bfloat16,
"param_dtype": jnp.bfloat16,
"input_shape": (1, 1)
}
train_arguments = TrainArguments(
model_class=type(model),
model_name="SFT-EasyDeL",
num_train_epochs=3,
configs_to_initialize_model_class=configs_to_initialize_model_class,
learning_rate=5e-5,
learning_rate_end=1e-6,
optimizer=EasyDelOptimizers.ADAMW,
scheduler=EasyDelSchedulers.WARM_UP_COSINE,
weight_decay=0.01,
total_batch_size=32,
max_training_steps=None, # None to let trainer Decide
do_train=True,
do_eval=False, # it's optional but supported
backend="tpu", # default backed is set to cpu, so you must define you want to use tpu cpu or gpu
max_length=max_length, # Note that you have to change this in the model config too
gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,
sharding_array=(1, -1, 1, 1), # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)
# everything training will be in sequence and model parallel automatic and share data between devices
remove_ckpt_after_load=True,
gradient_accumulation_steps=8,
loss_re_mat="",
dtype=jnp.bfloat16
)
def prompter(sample):
return [conversations_formatting_function(tokenizer, messages_field="messages")(sample)]
train_dataset = load_dataset("HuggingFaceH4/deita-10k-v0-sft", split="train_sft")
trainer = SFTTrainer(
arguments=train_arguments,
train_dataset=train_dataset,
eval_dataset=None, # we don't have eval dataset rn :)
tokenizer=tokenizer,
dataset_text_field=None,
formatting_func=prompter,
packing=False,
num_of_sequences=1024,
)
output = trainer.train(flax.core.FrozenDict({"params": params}))
print(f"Hey ! , here's where your model saved {output.checkpoint_path}")
```
> [!NOTE]
> You Can use Lora too, both for DPO and SFT Trainers.
## FineTuning
with using EasyDel FineTuning LLM (CausalLanguageModels) are easy as much as possible with using Jax and Flax
and having the benefit of TPUs for the best speed here's a simple code to use in order to finetune your
own Model
Days Has Been Passed and now using easydel in Jax is way more similar to HF/PyTorch Style
now it's time to finetune our model
```python
from EasyDel import (
TrainArguments,
CausalLanguageModelTrainer,
AutoEasyDelModelForCausalLM,
EasyDelOptimizers,
EasyDelSchedulers,
EasyDelGradientCheckPointers
)
from datasets import load_dataset
import flax
from jax import numpy as jnp
from transformers import AutoTokenizer
huggingface_repo_id_or_path = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
model, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )
max_length = 2048
tokenizer = AutoTokenizer.from_pretrained(
huggingface_repo_id_or_path,
trust_remote_code=True
)
tokenizer.pad_token = tokenizer.eos_token
configs_to_initialize_model_class = {
"config": model.config,
"dtype": jnp.bfloat16,
"param_dtype": jnp.bfloat16,
"input_shape": (1, 1)
}
train_arguments = TrainArguments(
model_class=type(model),
model_name="my_first_model_to_train_using_easydel",
num_train_epochs=3,
configs_to_initialize_model_class=configs_to_initialize_model_class,
learning_rate=5e-5,
learning_rate_end=1e-6,
optimizer=EasyDelOptimizers.ADAMW, # "adamw", "lion", "adafactor" are supported
scheduler=EasyDelSchedulers.LINEAR,
# "linear","cosine", "none" ,"warm_up_cosine" and "warm_up_linear" are supported
weight_decay=0.01,
total_batch_size=64,
max_training_steps=None, # None to let trainer Decide
do_train=True,
do_eval=False, # it's optional but supported
backend="tpu", # default backed is set to cpu, so you must define you want to use tpu cpu or gpu
max_length=max_length, # Note that you have to change this in the model config too
gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,
sharding_array=(1, -1, 1, 1), # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)
# everything training will be in sequence and model parallel automatic and share data between devices
remove_ckpt_after_load=True,
gradient_accumulation_steps=8,
loss_re_mat="",
dtype=jnp.bfloat16
)
def ultra_chat_prompting_process(
data_chunk
):
user_part = [
chunk["content"] for chunk in data_chunk["messages"] if chunk["role"] == "user"
]
assistant_part = [
chunk["content"] for chunk in data_chunk["messages"] if chunk["role"] == "assistant"
]
prompt = ""
for uc, ac in zip(user_part, assistant_part):
prompt += f"<|user|>\n{uc}</s>\n<|assistant|>\n{ac}</s>\n"
return {"prompt": prompt}
tokenization_process = lambda data_chunk: tokenizer(
data_chunk["prompt"],
add_special_tokens=False,
max_length=max_length,
padding="max_length"
)
dataset = load_dataset("HuggingFaceH4/ultrachat_200k")
dataset_train = dataset["train_gen"].map(ultra_chat_prompting_process, num_proc=12)
dataset_train = dataset_train.map(
tokenization_process,
num_proc=12,
remove_columns=dataset_train.column_names
)
# you can do the same for evaluation process dataset
trainer = CausalLanguageModelTrainer(
train_arguments,
dataset_train,
checkpoint_path=None
)
output = trainer.train(flax.core.FrozenDict({"params": params}))
print(f"Hey ! , here's where your model saved {output.checkpoint_path}")
```
> [!TIP]
> you can then convert it to pytorch for better use I don't recommend jax/flax for hosting models since
> pytorch is better option for gpus
## LLMServe
To use EasyDeL in your project, you will need to import the library in your Python script and use its various functions
and classes. Here is an example of how to import EasyDeL and use its Model class:
```python
from EasyDel.modules import AutoEasyDelModelForCausalLM
from EasyDel.serve import JAXServer
from transformers import AutoTokenizer
import jax
model_huggingface_repo_id = "meta-llama/Llama.md-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_huggingface_repo_id, trust_remote_code=True)
model, params = AutoEasyDelModelForCausalLM.from_pretrained(
model_huggingface_repo_id,
jax.devices("cpu")[0],
jax.numpy.float16,
jax.numpy.float16,
jax.lax.Precision("fastest"),
(1, -1, 1, 1),
device_map="auto"
)
server = JAXServer.from_parameters(
model=model,
config_model=model.config,
tokenizer=tokenizer,
params=model.params,
add_params_field=True
)
response_printed = 0
for response, tokens_used in server.process(
"String To The Model", stream=True
):
print(response[response_printed:], end="")
response_printed = len(response)
```
## DPO Fine-tuning
`DPOTrainer` is the new Trainer in EasyDeL, so you might have except some bugs in process but as far as i have tested
everything works just fine, and you can consider it the first DPO Trainer in JAX/Flax let have an example and see how
you can fine-tune your own model with DPOTrainer
> [!TIP]
> In case that you want a better script to learn about `DPOTrainer` you can see examples
> at [here](https://github.com/erfanzar/EasyDeL/blob/main/examples/training/dpo/dpo_training_example.py) which contain
> DPO Tuning a Mixtral model with Intel DPO dataset.
```python
from EasyDel import (
TrainArguments,
EasyDelOptimizers,
EasyDelSchedulers,
EasyDelGradientCheckPointers,
DPOTrainer,
EasyDelState,
easystate_to_huggingface_model
)
from datasets import load_dataset
from huggingface_hub import HfApi
from transformers import AutoTokenizer, LlamaForCausalLM as module_pt
from jax import numpy as jnp
import jax
from jax.sharding import PartitionSpec
from fjformer import GenerateRNG
from typing import Optional, Dict
from datasets import Dataset
rng_g = GenerateRNG()
api = HfApi()
max_length = 512 # Overall maximum length
max_target_length = 1024 # Maximum Length for target column in Dataset
max_prompt_length = 1024 # Maximum Length for prompt column in Dataset
model_name_or_path = "erfanzar/LinguaMatic-Tiny"
ref_model_name_or_path = "teknium/OpenHermes-2.5-Mistral-7B"
dtype = jnp.bfloat16
sharding_axis_dims = (1, -1, 1, 1)
sharding_axis_names = ("dp", "fsdp", "tp", "sp")
query_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Query Partition Spec for Model
key_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Key Partition Spec for Model
value_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Value Partition Spec for Model
bias_partition_spec = PartitionSpec(
("dp", "fsdp"), None, None, None
) # Attention Mask / Bias Partition Spec for Model
attention_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Attention Score / Weight Partition Spec for Model
ref_model_query_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Query Partition Spec for Ref Model
ref_model_key_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Key Partition Spec for Ref Model
ref_model_value_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Value Partition Spec for Ref Model
ref_model_bias_partition_spec = PartitionSpec(
("dp", "fsdp"), None, None, None
) # Attention Mask / Bias Partition Spec for Ref Model
ref_model_attention_partition_spec = PartitionSpec(
("dp", "fsdp"), "sp", "tp", None
) # Attention Score / Weight Partition Spec for Ref Model
def extract_anthropic_prompt(prompt_and_response):
"""
Extract the anthropic prompt from a prompt and response pair.
"""
search_term = "\n\nAssistant:"
search_term_idx = prompt_and_response.rfind(search_term)
assert search_term_idx != -1, f"Prompt and response does not contain '{search_term}'"
return prompt_and_response[: search_term_idx + len(search_term)]
def get_hh(split: str, sanity_check: bool = False, silent: bool = False, cache_dir: Optional[str] = None) -> Dataset:
"""
Load the Anthropic Helpful-Harmless dataset from Hugging Face and convert it to the necessary format.
The dataset is converted to a dictionary with the following structure:
{
'prompt': List[str],
'chosen': List[str],
'rejected': List[str],
}
Prompts should be structured as follows:
\n\nHuman: <prompt>\n\nAssistant:
Multiple turns are allowed, but the prompt should always start with \n\nHuman: and end with \n\nAssistant:.
"""
dataset = load_dataset("Anthropic/hh-rlhf", split=split, cache_dir=cache_dir)
if sanity_check:
dataset = dataset.select(range(min(len(dataset), 1000)))
def split_prompt_and_responses(sample) -> Dict[str, str]:
prompt = extract_anthropic_prompt(sample["chosen"])
return {
"prompt": prompt,
"chosen": sample["chosen"][len(prompt):],
"rejected": sample["rejected"][len(prompt):],
}
return dataset.map(split_prompt_and_responses)
arguments = TrainArguments(
model_name="EasyDeL-DPO",
num_train_epochs=5,
learning_rate=1e-4,
learning_rate_end=3e-5,
warmup_steps=200,
optimizer=EasyDelOptimizers.ADAMW,
scheduler=EasyDelSchedulers.LINEAR,
weight_decay=0.02,
total_batch_size=128,
gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,
sharding_array=sharding_axis_dims,
fully_sharded_data_parallel=True,
gradient_accumulation_steps=2,
dtype=dtype,
param_dtype=dtype,
step_start_point=0,
training_time="7H",
do_train=True,
do_eval=True,
track_memory=False # Performance boost.
# You can set other options too or play with them but for now I just stick with these arguments.
)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
train_dataset = get_hh("train", sanity_check=True)
eval_dataset = get_hh("test", sanity_check=True)
state = EasyDelState.from_pretrained(
pretrained_model_name_or_path=model_name_or_path,
dtype=dtype,
param_dtype=dtype,
init_optimizer_state=False,
free_optimizer_state=True,
sharding_axis_dims=sharding_axis_dims,
sharding_axis_names=sharding_axis_names,
query_partition_spec=query_partition_spec,
key_partition_spec=key_partition_spec,
value_partition_spec=value_partition_spec,
bias_partition_spec=bias_partition_spec,
attention_partition_spec=attention_partition_spec,
)
ref_state = EasyDelState.from_pretrained(
pretrained_model_name_or_path=ref_model_name_or_path,
dtype=dtype,
param_dtype=dtype,
init_optimizer_state=False,
free_optimizer_state=True,
sharding_axis_dims=sharding_axis_dims,
sharding_axis_names=sharding_axis_names,
query_partition_spec=ref_model_query_partition_spec,
key_partition_spec=ref_model_key_partition_spec,
value_partition_spec=ref_model_value_partition_spec,
bias_partition_spec=ref_model_bias_partition_spec,
attention_partition_spec=ref_model_attention_partition_spec,
)
dpo_trainer = DPOTrainer(
model_state=state,
ref_model_state=ref_state,
beta=0.1,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
arguments=arguments,
max_length=max_length,
max_target_length=max_target_length,
max_prompt_length=max_prompt_length,
ref_model_init_kwargs=None, # In case that you pass the ref_model_state a string you have to pass this one too
model_init_kwargs=None, # In case that you pass the model_state a string you have to pass this one too
dataset_map_arguments={
"num_proc": 8,
"batched": True,
"batch_size": 100,
},
auto_shard_model_state=True,
auto_shard_ref_model_state=True,
loss_type="sigmoid",
data_collator=None, # Pass None in order to use default data_collector (you can create your own)
)
output = dpo_trainer.train()
easydel_jax_model = output.state # Here's you EasyDeL Model
with jax.default_device(jax.devices("cpu")[0]):
model = easystate_to_huggingface_model(
state=EasyDelState.load_state(
output.checkpoint_path
),
base_huggingface_module=module_pt,
config=dpo_trainer.model_state.module.config
) # Here's you PyTorch Model
model.push_to_hub("<REPO_ID>", private=False) # Hope you love open-source too :)
tokenizer.push_to_hub("<REPO_ID>", private=False) # Hope you love open-source too :)
```
now you have trained your first model Using DPOTrainer in JAX with EasyDeL.
> [!TIP]
> The API of EasyDeL DPO Trainer is similar to DPO Trainer in TRL from HuggingFace so that means
> you have freedom and have access to a hackable and changeable code.
## EasyDelState
EasyDelState is new and cool feature in EasyDeL and have a lot of options like
storing `Model Parameters`, _Optimizer State,
Model Config, Model Type, Optimizer and Scheduler Configs_
Let see and examples of using EasyDelState
### Fine-tuning
Fine-tuning from a previous State or a new state
```python
from EasyDel import (
AutoEasyDelConfig,
EasyDelState
)
from transformers import AutoTokenizer
from jax import numpy as jnp, lax
import jax
huggingface_model_repo_id = "REPO_ID"
checkpoint_name = "CKPT_NAME"
state = EasyDelState.from_pretrained(
pretrained_model_name_or_path=huggingface_model_repo_id,
filename=checkpoint_name,
optimizer="adamw",
scheduler="none",
tx_init=None,
device=jax.devices('cpu')[0], # Offload Device
dtype=jnp.bfloat16,
param_dtype=jnp.bfloat16,
precision=lax.Precision("fastest"),
sharding_axis_dims=(1, -1, 1, 1),
# the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)
# everything training will be in sequence and model parallel automatic and share data between devices
sharding_axis_names=("dp", "fsdp", "tp", "sp"),
query_partition_spec=jax.sharding.PartitionSpec(("dp", "fsdp"), "sp", "tp", None),
key_partition_spec=jax.sharding.PartitionSpec(("dp", "fsdp"), "sp", "tp", None),
value_partition_spec=jax.sharding.PartitionSpec(("dp", "fsdp"), "sp", "tp", None),
bias_partition_spec=jax.sharding.PartitionSpec(("dp", "fsdp"), None, None, None),
attention_partition_spec=jax.sharding.PartitionSpec(("dp", "fsdp"), "sp", "tp", None),
shard_attention_computation=True,
input_shape=(1, 1),
backend=None,
init_optimizer_state=False,
free_optimizer_state=True,
verbose=True,
state_shard_fns=None,
)
config = AutoEasyDelConfig.from_pretrained(
huggingface_model_repo_id
)
tokenizer = AutoTokenizer.from_pretrained(
huggingface_model_repo_id,
trust_remote_code=True
)
max_length = config.max_position_embeddings
configs_to_initialize_model_class = {
'config': config,
'dtype': jnp.bfloat16,
'param_dtype': jnp.bfloat16,
'input_shape': (8, 8)
}
```
`EasyDelState` also has `.load_state()` and `.save_state()` with some other usable options like `.free_opt_state()`
which
free optimizer state or `.shard_params()` which shard parameters you can read docs in order to find out more about these
options.
### Converting to Huggingface and Pytorch
Let see how you can convert a EasyDelMistral Model to Huggingface Pytorch Mistral Model from a trained State
```python
from transformers import MistralForCausalLM
from EasyDel import (
AutoEasyDelConfig,
EasyDelState,
easystate_to_huggingface_model
)
import jax
huggingface_model_repo_id = "REPO_ID"
config = AutoEasyDelConfig.from_pretrained(
huggingface_model_repo_id
)
with jax.default_device(jax.devices("cpu")[0]):
model = easystate_to_huggingface_model(
state=EasyDelState.load_state(
"PATH_TO_CKPT"
), # You can Pass EasyDelState here
base_huggingface_module=MistralForCausalLM, # type: ignore
config=config
)
model = model.half() # it's a huggingface model now
```
### Other Use Cases
`EasyDelState` have a general use you can use it everywhere in easydel for example for a stand-alone model
, serve, fine-tuning and many other features, it's up to you to test how creative you are 😇.
## Flash Attention and Splash Attention Are Here 🥵
here's a simple example about how can you use Flash Attention in EasyDeL
```python
# Config is built in config for every model (EasyDelPretrainedConfig)
config.add_basic_configurations(
attn_mechanism="flash", # Any supported Attention Mechanism
block_b=1,
block_q=512,
block_k=512,
block_k_major=512
)
```
_Flash Attention works on TPU with ease but for gpu there are still some improvements in process._
> [!TIP]
> use these partition specs in case of not using custom sharding_axis_names and using sequence sharding with flash
> flash attention
> ```python
>query_partition_spec=PartitionSpec(("dp", "fsdp"), None, "sp", "tp"),
>generation_query_partition_spec=PartitionSpec(("dp", "fsdp"), None, None, "tp"),
>key_partition_spec=PartitionSpec(("dp", "fsdp"), None, "sp", "tp"),
>value_partition_spec=PartitionSpec(("dp", "fsdp"), None, "sp", "tp"),
>attention_partition_spec=PartitionSpec(("dp", "fsdp"), None,"sp", "tp"),
> ```
## EasyDeLXRapTure for layer tuning and LoRA
in case of using LoRA and applying that on the EasyDeL models there are some other things
that you might need to config on your own but a lot of things being handled by EasyDeL so let just jump into an example
for LoRA fine-tuning section and use _EasyDeLXRapTure_ in for mistral models with flash attention example
```python
from flax.core import FrozenDict
from EasyDel import (
TrainArguments,
CausalLanguageModelTrainer,
AutoEasyDelModelForCausalLM,
EasyDelOptimizers,
EasyDelSchedulers,
EasyDelGradientCheckPointers,
EasyDeLXRapTureConfig
)
from datasets import load_dataset
import flax
from jax import numpy as jnp
from transformers import AutoTokenizer
huggingface_repo_id_or_path = "mistralai/Mistral-7B-Instruct-v0.1"
model, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )
max_length = 8196
model_parameters = FrozenDict({"params": params})
dtype = jnp.bfloat16
param_dtype = jnp.bfloat16 # you can change that if you want
tokenizer = AutoTokenizer.from_pretrained(
huggingface_repo_id_or_path,
trust_remote_code=True
)
model.config.add_basic_configurations(
attn_mechanism="flash", # Using FlashAttention
block_b=1,
block_q=1024,
block_k=1024,
block_k_major=1024,
)
tokenizer.pad_token = tokenizer.eos_token
configs_to_initialize_model_class = {
"config": model.config,
"dtype": dtype,
"param_dtype": param_dtype,
"input_shape": (1, 1)
}
rapture = EasyDeLXRapTureConfig(
parameters=model_parameters,
lora_dim=64,
fully_fine_tune_parameters=["embed_tokens"], # Model layer to be fully fine tuned
lora_fine_tune_parameters=["q_proj", "v_proj", "k_proj", "o_proj"], # LoRA Layer Targets you can pass this to none
# For only Layer Tuning or transfer learning
verbose=True
)
train_arguments = TrainArguments(
model_class=type(model),
model_name="EasyDeL-Lora-Example",
num_train_epochs=3,
configs_to_initialize_model_class=configs_to_initialize_model_class,
learning_rate=1e-4, # Using higher learning rate is recommended
learning_rate_end=8e-5,
optimizer=EasyDelOptimizers.ADAMW, # "adamw", "lion", "adafactor" are supported
scheduler=EasyDelSchedulers.LINEAR,
# "linear","cosine", "none" ,"warm_up_cosine" and "warm_up_linear" are supported
weight_decay=0.01,
total_batch_size=512,
max_training_steps=None, # None to let trainer Decide
do_train=True,
do_eval=False, # it's optional but supported
backend="tpu", # default backed is set to cpu, so you must define you want to use tpu cpu or gpu
max_length=max_length, # Note that you have to change this in the model config too
gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,
sharding_array=(1, -1, 1, 1), # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)
# everything training will be in sequence and model parallel automatic and share data between devices
remove_ckpt_after_load=True,
gradient_accumulation_steps=1,
loss_re_mat="",
dtype=dtype,
param_dtype=param_dtype,
rapture_config=rapture,
merge_lora_rapture_parameters=True # turning this off is still not supported and not recommended to do so
# What this does ? this will merge the lora parameters with the original model parameters and the end of training
)
def ultra_chat_prompting_process(
data_chunk
):
user_part = [
chunk["content"] for chunk in data_chunk["messages"] if chunk["role"] == "user"
]
assistant_part = [
chunk["content"] for chunk in data_chunk["messages"] if chunk["role"] == "assistant"
]
prompt = ""
for uc, ac in zip(user_part, assistant_part):
prompt += f"<|user|>\n{uc}</s>\n<|assistant|>\n{ac}</s>\n"
return {"prompt": prompt}
tokenization_process = lambda data_chunk: tokenizer(
data_chunk["prompt"],
add_special_tokens=False,
max_length=max_length,
padding="max_length"
)
dataset = load_dataset("HuggingFaceH4/ultrachat_200k")
dataset_train = dataset["train_gen"].map(ultra_chat_prompting_process, num_proc=12)
dataset_train = dataset_train.map(
tokenization_process,
num_proc=12,
remove_columns=dataset_train.column_names
)
# you can do the same for evaluation process dataset
trainer = CausalLanguageModelTrainer(
train_arguments,
dataset_train,
checkpoint_path=None
)
output = trainer.train() # you should not pass the parameters in Trainer.train anymore when
# you are using LoRA or transfer Learning
print(f"Hey ! , here's where your model saved {output.checkpoint_path}")
```
## Contributing
EasyDeL is an open-source project, and contributions are welcome. If you would like to contribute to EasyDeL, please
fork the repository, make your changes, and submit a pull request. The team behind EasyDeL will review your changes and
merge them if they are suitable.
## License 📜
EasyDeL is an Fully Open-Source released under the Apache v2 license. Please see the LICENSE file in the root directory
of this project for
more information.
## Contact
If you have any questions or comments about EasyDeL, you can reach out to me
## Citing EasyDeL 🥶
To cite this repository:
```misc
@misc{Zare Chavoshi_2023,
title={EasyDeL, an open-source library, is specifically designed to enhance and streamline the training process of machine learning models. It focuses primarily on Jax/Flax and aims to provide convenient and effective solutions for training Flax/Jax Models on TPU/GPU for both Serving and Training purposes.},
url={https://github.com/erfanzar/EasyDel},
journal={EasyDeL Easy and Fast DeepLearning with JAX},
publisher={Erfan Zare Chavoshi},
author={Zare Chavoshi, Erfan},
year={2023}
}
```
Raw data
{
"_id": null,
"home_page": null,
"name": "EasyDeL",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.8",
"maintainer_email": null,
"keywords": "JAX, Torch, Deep Learning, Machine Learning, Flax, XLA",
"author": null,
"author_email": "Erfan Zare Chavoshi <Erfanzare810@gmail.com>",
"download_url": "https://files.pythonhosted.org/packages/7f/86/87337a43b7fd5a984085e7c0c34683db860c61875fec900f66dc271bea28/easydel-0.0.63.tar.gz",
"platform": null,
"description": "# EasyDeL \ud83d\udd2e\n\nEasyDeL, an open-source library, is specifically designed to enhance and streamline the training process of machine\nlearning models. It focuses primarily on Jax/Flax and aims to provide convenient and effective solutions for training\nFlax/Jax Models on TPU/GPU for both Serving and Training purposes. Additionally, EasyDeL will support mojo and will be\nrewritten for mojo as well.\n\nSome of the key features provided by EasyDeL include:\n\n- DPOTrainer, SFTTrainer, and VideoCLM Trainers\n- Serving and API Engines for Using and serving LLMs in JAX as efficiently as possible.\n- Support Quantization Methods for all the Models.\n- Support for 8, 6, and 4 BIT Operation, for inference and training in JAX\n- A wide range of models in Jax is supported which have never been implemented before such as Falcon, Qwen2, Phi2,\n Mixtral, Qwen2Moe, Cohere,and MPT ...\n- Integration of flashAttention in JAX for GPUs and TPUs\n- Automatic serving of LLMs with mid and high-level APIs in both JAX and PyTorch\n- LLM Trainer and fine-tuner in JAX\n- Video CLM Trainer and Fine-tuner for Models such Falcon, Qwen2, Phi2, MPT, Mixtral, Grok-1, and Qwen2Moe ...\n- RLHF (Reinforcement Learning from Human Feedback) in Jax (Beta Stage)\n- Various other features to enhance the training process and optimize performance.\n- LoRA: Low-Rank Adaptation of Large Language Models\n- RingAttention, Flash Attention, BlockWise FFN, and Efficient Attention are supported for more than 90 % of models\n ([FJFormer](https://github.com/erfanzar/FJFormer) Backbone).\n- Serving and API Engines for Using and serving LLMs in JAX as efficient as possible.\n- Automatic Converting Models from JAX-EasyDeL to PyTorch-HF and reverse\n\n> **News**\n> \n> Phi3 Model bugs are fixed, Arctic Model is added.\n>\n> Phi3 Model is present Now Apr 24 2024\n>\n> Drbx Model is present\n> Now [Apr 23 2024](https://github.com/erfanzar/EasyDeL/commit/4c1c5af099dad9334a82808eb04b03e7e567ddb7)\n>\n> New Attention Types are added `sharded_vanilla`, `wise_ring`, `sharded_vanilla` is same as `vanilla` but\n> uses shard_map which will make it a little faster and more efficent.\n>\n> `local_ring` is Added which is Ring Attention but for TPU/GPU/CPU(s) and support attention bias instead of attention\n> mask, `normal` attention is renamed to `vanilla`.\n>\n> `load_in_8bit` is now available for all the models, and requires to upgrade _fjformer to 0.0.50_\n>\n> Sharing Key and Value Cache for Large Sequence Length across devices are now Fixed (Attention Models).\n>\n> Cohere Model added [Apr 14 2024](https://github.com/erfanzar/EasyDeL/commit/06acb4a1afd7b67982a88b50840b90e73b1c9850)\n\n## Documentation \ud83d\udcab\n\n> [!IMPORTANT]\n> Documents and Examples are ready at [Here](https://erfanzar.github.io/EasyDeL)\n> Please have that in mind that EasyDel is in the loop of fast-development\n> so we might have API changes.\n\n### Hands on Code Kaggle Examples\n\n1. [script](https://www.kaggle.com/citifer/easydel-causal-language-model-trainer-example) for mindset of using EasyDeL\n CausalLanguageModelTrainer on kaggle, but you can do much more.\n2. [script](https://www.kaggle.com/code/citifer/easydel-serve-example-mixtral) for using and serving LLMs with EasyDeL\n JAXServer API (Mixtral Example).\n3. [script](https://www.kaggle.com/code/citifer/easydel-sfttrainer-example) SuperVised Finetuning with EasyDeL.\n\n## Serving\n\nyou can read docs or examples to see how `JAXServer` works but let me show you how you can simply host and serve any\nmodel that supported by `EasyDeL` fo this example ill just use `gemma-7-it` by google, but you can use any model as you\nwish.\n\n```shell\npython -m examples.jax_serve_example \\\n --prompter_type=\"gemma\" \\ \n --share_gradio=True \\ \n --sharding_axis_dims=1,1,1,-1 \\\n --attn_mechanism=\"sharded_vanilla\" \\\n --scan_ring_attention=True \\\n --max_sequence_length=8192 \\ \n --max_new_tokens_ratio=25 \\\n --max_compile_tokens=256 \\ \n --block_k=128 \\\n --block_q=128 \\\n --pretrained_model_name_or_path=\"google/gemma-7b-it\" \\\n --dtype=\"bf16\"\n```\n\n> [!NOTE]\n> you can use `EasyServe` which is a Serve API Engine for production purpose sice that's more stable provide versioned\n> API and efficient.\n\n## Supervised Fine-Tuning with EasyDeL\n\nEasyDeL supports both DPO and SFT Trainers, so dealing with LLMs in jax is a lot easier right now\nlet have an example of using Supervised Fine-Tuner in JAX with EasyDeL\n\n```python\nfrom EasyDel import (\n TrainArguments,\n AutoEasyDelModelForCausalLM,\n EasyDelOptimizers,\n EasyDelSchedulers,\n EasyDelGradientCheckPointers,\n SFTTrainer,\n conversations_formatting_function # i have added this one for newcomers so if they \n # don't know what's going on they can use this pre created prompter\n)\nfrom datasets import load_dataset\nimport flax\nfrom jax import numpy as jnp\nfrom transformers import AutoTokenizer\n\nhuggingface_repo_id_or_path = \"mistralai/Mistral-7B-Instruct-v0.2\"\n\nmodel, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )\n\nmax_length = 4096\ntokenizer = AutoTokenizer.from_pretrained(\n huggingface_repo_id_or_path,\n trust_remote_code=True\n)\ntokenizer.pad_token = tokenizer.eos_token\nconfigs_to_initialize_model_class = {\n \"config\": model.config,\n \"dtype\": jnp.bfloat16,\n \"param_dtype\": jnp.bfloat16,\n \"input_shape\": (1, 1)\n}\n\ntrain_arguments = TrainArguments(\n model_class=type(model),\n model_name=\"SFT-EasyDeL\",\n num_train_epochs=3,\n configs_to_initialize_model_class=configs_to_initialize_model_class,\n learning_rate=5e-5,\n learning_rate_end=1e-6,\n optimizer=EasyDelOptimizers.ADAMW,\n scheduler=EasyDelSchedulers.WARM_UP_COSINE,\n weight_decay=0.01,\n total_batch_size=32,\n max_training_steps=None, # None to let trainer Decide\n do_train=True,\n do_eval=False, # it's optional but supported \n backend=\"tpu\", # default backed is set to cpu, so you must define you want to use tpu cpu or gpu\n max_length=max_length, # Note that you have to change this in the model config too\n gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,\n sharding_array=(1, -1, 1, 1), # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)\n # everything training will be in sequence and model parallel automatic and share data between devices\n remove_ckpt_after_load=True,\n gradient_accumulation_steps=8,\n loss_re_mat=\"\",\n dtype=jnp.bfloat16\n)\n\n\ndef prompter(sample):\n return [conversations_formatting_function(tokenizer, messages_field=\"messages\")(sample)]\n\n\ntrain_dataset = load_dataset(\"HuggingFaceH4/deita-10k-v0-sft\", split=\"train_sft\")\ntrainer = SFTTrainer(\n arguments=train_arguments,\n train_dataset=train_dataset,\n eval_dataset=None, # we don't have eval dataset rn :)\n tokenizer=tokenizer,\n dataset_text_field=None,\n\n formatting_func=prompter,\n packing=False,\n num_of_sequences=1024,\n)\n\noutput = trainer.train(flax.core.FrozenDict({\"params\": params}))\nprint(f\"Hey ! , here's where your model saved {output.checkpoint_path}\")\n```\n\n> [!NOTE]\n> You Can use Lora too, both for DPO and SFT Trainers.\n\n## FineTuning\n\nwith using EasyDel FineTuning LLM (CausalLanguageModels) are easy as much as possible with using Jax and Flax\nand having the benefit of TPUs for the best speed here's a simple code to use in order to finetune your\nown Model\n\nDays Has Been Passed and now using easydel in Jax is way more similar to HF/PyTorch Style\nnow it's time to finetune our model\n\n```python\nfrom EasyDel import (\n TrainArguments,\n CausalLanguageModelTrainer,\n AutoEasyDelModelForCausalLM,\n EasyDelOptimizers,\n EasyDelSchedulers,\n EasyDelGradientCheckPointers\n)\nfrom datasets import load_dataset\nimport flax\nfrom jax import numpy as jnp\nfrom transformers import AutoTokenizer\n\nhuggingface_repo_id_or_path = \"TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\"\n\nmodel, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )\n\nmax_length = 2048\ntokenizer = AutoTokenizer.from_pretrained(\n huggingface_repo_id_or_path,\n trust_remote_code=True\n)\ntokenizer.pad_token = tokenizer.eos_token\nconfigs_to_initialize_model_class = {\n \"config\": model.config,\n \"dtype\": jnp.bfloat16,\n \"param_dtype\": jnp.bfloat16,\n \"input_shape\": (1, 1)\n}\n\ntrain_arguments = TrainArguments(\n model_class=type(model),\n model_name=\"my_first_model_to_train_using_easydel\",\n num_train_epochs=3,\n configs_to_initialize_model_class=configs_to_initialize_model_class,\n learning_rate=5e-5,\n learning_rate_end=1e-6,\n optimizer=EasyDelOptimizers.ADAMW, # \"adamw\", \"lion\", \"adafactor\" are supported\n scheduler=EasyDelSchedulers.LINEAR,\n # \"linear\",\"cosine\", \"none\" ,\"warm_up_cosine\" and \"warm_up_linear\" are supported\n weight_decay=0.01,\n total_batch_size=64,\n max_training_steps=None, # None to let trainer Decide\n do_train=True,\n do_eval=False, # it's optional but supported \n backend=\"tpu\", # default backed is set to cpu, so you must define you want to use tpu cpu or gpu\n max_length=max_length, # Note that you have to change this in the model config too\n gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,\n sharding_array=(1, -1, 1, 1), # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)\n # everything training will be in sequence and model parallel automatic and share data between devices\n remove_ckpt_after_load=True,\n gradient_accumulation_steps=8,\n loss_re_mat=\"\",\n dtype=jnp.bfloat16\n)\n\n\ndef ultra_chat_prompting_process(\n data_chunk\n):\n user_part = [\n chunk[\"content\"] for chunk in data_chunk[\"messages\"] if chunk[\"role\"] == \"user\"\n ]\n assistant_part = [\n chunk[\"content\"] for chunk in data_chunk[\"messages\"] if chunk[\"role\"] == \"assistant\"\n ]\n\n prompt = \"\"\n\n for uc, ac in zip(user_part, assistant_part):\n prompt += f\"<|user|>\\n{uc}</s>\\n<|assistant|>\\n{ac}</s>\\n\"\n\n return {\"prompt\": prompt}\n\n\ntokenization_process = lambda data_chunk: tokenizer(\n data_chunk[\"prompt\"],\n add_special_tokens=False,\n max_length=max_length,\n padding=\"max_length\"\n)\n\ndataset = load_dataset(\"HuggingFaceH4/ultrachat_200k\")\ndataset_train = dataset[\"train_gen\"].map(ultra_chat_prompting_process, num_proc=12)\ndataset_train = dataset_train.map(\n tokenization_process,\n num_proc=12,\n remove_columns=dataset_train.column_names\n)\n\n# you can do the same for evaluation process dataset\n\ntrainer = CausalLanguageModelTrainer(\n train_arguments,\n dataset_train,\n checkpoint_path=None\n)\n\noutput = trainer.train(flax.core.FrozenDict({\"params\": params}))\nprint(f\"Hey ! , here's where your model saved {output.checkpoint_path}\")\n```\n\n> [!TIP]\n> you can then convert it to pytorch for better use I don't recommend jax/flax for hosting models since\n> pytorch is better option for gpus\n\n## LLMServe\n\nTo use EasyDeL in your project, you will need to import the library in your Python script and use its various functions\nand classes. Here is an example of how to import EasyDeL and use its Model class:\n\n```python\nfrom EasyDel.modules import AutoEasyDelModelForCausalLM\nfrom EasyDel.serve import JAXServer\nfrom transformers import AutoTokenizer\nimport jax\n\nmodel_huggingface_repo_id = \"meta-llama/Llama.md-2-7b-chat-hf\"\n\ntokenizer = AutoTokenizer.from_pretrained(model_huggingface_repo_id, trust_remote_code=True)\nmodel, params = AutoEasyDelModelForCausalLM.from_pretrained(\n model_huggingface_repo_id,\n jax.devices(\"cpu\")[0],\n jax.numpy.float16,\n jax.numpy.float16,\n jax.lax.Precision(\"fastest\"),\n (1, -1, 1, 1),\n device_map=\"auto\"\n)\n\nserver = JAXServer.from_parameters(\n model=model,\n config_model=model.config,\n tokenizer=tokenizer,\n params=model.params,\n add_params_field=True\n)\n\nresponse_printed = 0\nfor response, tokens_used in server.process(\n \"String To The Model\", stream=True\n):\n print(response[response_printed:], end=\"\")\n response_printed = len(response)\n```\n\n## DPO Fine-tuning\n\n`DPOTrainer` is the new Trainer in EasyDeL, so you might have except some bugs in process but as far as i have tested\neverything works just fine, and you can consider it the first DPO Trainer in JAX/Flax let have an example and see how\nyou can fine-tune your own model with DPOTrainer\n\n> [!TIP]\n> In case that you want a better script to learn about `DPOTrainer` you can see examples\n> at [here](https://github.com/erfanzar/EasyDeL/blob/main/examples/training/dpo/dpo_training_example.py) which contain\n> DPO Tuning a Mixtral model with Intel DPO dataset.\n\n```python\nfrom EasyDel import (\n TrainArguments,\n EasyDelOptimizers,\n EasyDelSchedulers,\n EasyDelGradientCheckPointers,\n DPOTrainer,\n EasyDelState,\n easystate_to_huggingface_model\n)\n\nfrom datasets import load_dataset\nfrom huggingface_hub import HfApi\nfrom transformers import AutoTokenizer, LlamaForCausalLM as module_pt\nfrom jax import numpy as jnp\nimport jax\nfrom jax.sharding import PartitionSpec\nfrom fjformer import GenerateRNG\nfrom typing import Optional, Dict\nfrom datasets import Dataset\n\nrng_g = GenerateRNG()\napi = HfApi()\n\nmax_length = 512 # Overall maximum length\nmax_target_length = 1024 # Maximum Length for target column in Dataset\nmax_prompt_length = 1024 # Maximum Length for prompt column in Dataset\n\nmodel_name_or_path = \"erfanzar/LinguaMatic-Tiny\"\nref_model_name_or_path = \"teknium/OpenHermes-2.5-Mistral-7B\"\ndtype = jnp.bfloat16\n\nsharding_axis_dims = (1, -1, 1, 1)\nsharding_axis_names = (\"dp\", \"fsdp\", \"tp\", \"sp\")\nquery_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Query Partition Spec for Model\nkey_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Key Partition Spec for Model\nvalue_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Value Partition Spec for Model\nbias_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), None, None, None\n) # Attention Mask / Bias Partition Spec for Model\nattention_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Attention Score / Weight Partition Spec for Model\n\nref_model_query_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Query Partition Spec for Ref Model\nref_model_key_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Key Partition Spec for Ref Model\nref_model_value_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Value Partition Spec for Ref Model\nref_model_bias_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), None, None, None\n) # Attention Mask / Bias Partition Spec for Ref Model\nref_model_attention_partition_spec = PartitionSpec(\n (\"dp\", \"fsdp\"), \"sp\", \"tp\", None\n) # Attention Score / Weight Partition Spec for Ref Model\n\n\ndef extract_anthropic_prompt(prompt_and_response):\n \"\"\"\n Extract the anthropic prompt from a prompt and response pair.\n \"\"\"\n search_term = \"\\n\\nAssistant:\"\n search_term_idx = prompt_and_response.rfind(search_term)\n assert search_term_idx != -1, f\"Prompt and response does not contain '{search_term}'\"\n return prompt_and_response[: search_term_idx + len(search_term)]\n\n\ndef get_hh(split: str, sanity_check: bool = False, silent: bool = False, cache_dir: Optional[str] = None) -> Dataset:\n \"\"\"\n Load the Anthropic Helpful-Harmless dataset from Hugging Face and convert it to the necessary format.\n\n The dataset is converted to a dictionary with the following structure:\n {\n 'prompt': List[str],\n 'chosen': List[str],\n 'rejected': List[str],\n }\n\n Prompts should be structured as follows:\n \\n\\nHuman: <prompt>\\n\\nAssistant:\n Multiple turns are allowed, but the prompt should always start with \\n\\nHuman: and end with \\n\\nAssistant:.\n \"\"\"\n dataset = load_dataset(\"Anthropic/hh-rlhf\", split=split, cache_dir=cache_dir)\n if sanity_check:\n dataset = dataset.select(range(min(len(dataset), 1000)))\n\n def split_prompt_and_responses(sample) -> Dict[str, str]:\n prompt = extract_anthropic_prompt(sample[\"chosen\"])\n return {\n \"prompt\": prompt,\n \"chosen\": sample[\"chosen\"][len(prompt):],\n \"rejected\": sample[\"rejected\"][len(prompt):],\n }\n\n return dataset.map(split_prompt_and_responses)\n\n\narguments = TrainArguments(\n model_name=\"EasyDeL-DPO\",\n num_train_epochs=5,\n learning_rate=1e-4,\n learning_rate_end=3e-5,\n warmup_steps=200,\n optimizer=EasyDelOptimizers.ADAMW,\n scheduler=EasyDelSchedulers.LINEAR,\n weight_decay=0.02,\n total_batch_size=128,\n gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,\n sharding_array=sharding_axis_dims,\n fully_sharded_data_parallel=True,\n gradient_accumulation_steps=2,\n dtype=dtype,\n param_dtype=dtype,\n step_start_point=0,\n training_time=\"7H\",\n do_train=True,\n do_eval=True,\n track_memory=False # Performance boost.\n # You can set other options too or play with them but for now I just stick with these arguments.\n)\n\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n\nif tokenizer.pad_token is None:\n tokenizer.pad_token = tokenizer.eos_token\n\nif tokenizer.pad_token_id is None:\n tokenizer.pad_token_id = tokenizer.eos_token_id\n\ntrain_dataset = get_hh(\"train\", sanity_check=True)\neval_dataset = get_hh(\"test\", sanity_check=True)\n\nstate = EasyDelState.from_pretrained(\n pretrained_model_name_or_path=model_name_or_path,\n dtype=dtype,\n param_dtype=dtype,\n init_optimizer_state=False,\n free_optimizer_state=True,\n sharding_axis_dims=sharding_axis_dims,\n sharding_axis_names=sharding_axis_names,\n query_partition_spec=query_partition_spec,\n key_partition_spec=key_partition_spec,\n value_partition_spec=value_partition_spec,\n bias_partition_spec=bias_partition_spec,\n attention_partition_spec=attention_partition_spec,\n)\n\nref_state = EasyDelState.from_pretrained(\n pretrained_model_name_or_path=ref_model_name_or_path,\n dtype=dtype,\n param_dtype=dtype,\n init_optimizer_state=False,\n free_optimizer_state=True,\n sharding_axis_dims=sharding_axis_dims,\n sharding_axis_names=sharding_axis_names,\n query_partition_spec=ref_model_query_partition_spec,\n key_partition_spec=ref_model_key_partition_spec,\n value_partition_spec=ref_model_value_partition_spec,\n bias_partition_spec=ref_model_bias_partition_spec,\n attention_partition_spec=ref_model_attention_partition_spec,\n)\n\ndpo_trainer = DPOTrainer(\n model_state=state,\n ref_model_state=ref_state,\n beta=0.1,\n train_dataset=train_dataset,\n eval_dataset=eval_dataset,\n tokenizer=tokenizer,\n arguments=arguments,\n max_length=max_length,\n max_target_length=max_target_length,\n max_prompt_length=max_prompt_length,\n ref_model_init_kwargs=None, # In case that you pass the ref_model_state a string you have to pass this one too\n model_init_kwargs=None, # In case that you pass the model_state a string you have to pass this one too\n dataset_map_arguments={\n \"num_proc\": 8,\n \"batched\": True,\n \"batch_size\": 100,\n },\n auto_shard_model_state=True,\n auto_shard_ref_model_state=True,\n loss_type=\"sigmoid\",\n data_collator=None, # Pass None in order to use default data_collector (you can create your own)\n)\n\noutput = dpo_trainer.train()\n\neasydel_jax_model = output.state # Here's you EasyDeL Model\n\nwith jax.default_device(jax.devices(\"cpu\")[0]):\n model = easystate_to_huggingface_model(\n state=EasyDelState.load_state(\n output.checkpoint_path\n ),\n base_huggingface_module=module_pt,\n config=dpo_trainer.model_state.module.config\n ) # Here's you PyTorch Model\n\nmodel.push_to_hub(\"<REPO_ID>\", private=False) # Hope you love open-source too :)\ntokenizer.push_to_hub(\"<REPO_ID>\", private=False) # Hope you love open-source too :)\n```\n\nnow you have trained your first model Using DPOTrainer in JAX with EasyDeL.\n\n> [!TIP]\n> The API of EasyDeL DPO Trainer is similar to DPO Trainer in TRL from HuggingFace so that means\n> you have freedom and have access to a hackable and changeable code.\n\n## EasyDelState\n\nEasyDelState is new and cool feature in EasyDeL and have a lot of options like\nstoring `Model Parameters`, _Optimizer State,\nModel Config, Model Type, Optimizer and Scheduler Configs_\n\nLet see and examples of using EasyDelState\n\n### Fine-tuning\n\nFine-tuning from a previous State or a new state\n\n```python\nfrom EasyDel import (\n AutoEasyDelConfig,\n EasyDelState\n)\nfrom transformers import AutoTokenizer\nfrom jax import numpy as jnp, lax\nimport jax\n\nhuggingface_model_repo_id = \"REPO_ID\"\ncheckpoint_name = \"CKPT_NAME\"\n\nstate = EasyDelState.from_pretrained(\n pretrained_model_name_or_path=huggingface_model_repo_id,\n filename=checkpoint_name,\n optimizer=\"adamw\",\n scheduler=\"none\",\n tx_init=None,\n device=jax.devices('cpu')[0], # Offload Device\n dtype=jnp.bfloat16,\n param_dtype=jnp.bfloat16,\n precision=lax.Precision(\"fastest\"),\n sharding_axis_dims=(1, -1, 1, 1),\n # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)\n # everything training will be in sequence and model parallel automatic and share data between devices\n sharding_axis_names=(\"dp\", \"fsdp\", \"tp\", \"sp\"),\n query_partition_spec=jax.sharding.PartitionSpec((\"dp\", \"fsdp\"), \"sp\", \"tp\", None),\n key_partition_spec=jax.sharding.PartitionSpec((\"dp\", \"fsdp\"), \"sp\", \"tp\", None),\n value_partition_spec=jax.sharding.PartitionSpec((\"dp\", \"fsdp\"), \"sp\", \"tp\", None),\n bias_partition_spec=jax.sharding.PartitionSpec((\"dp\", \"fsdp\"), None, None, None),\n attention_partition_spec=jax.sharding.PartitionSpec((\"dp\", \"fsdp\"), \"sp\", \"tp\", None),\n shard_attention_computation=True,\n input_shape=(1, 1),\n backend=None,\n init_optimizer_state=False,\n free_optimizer_state=True,\n verbose=True,\n state_shard_fns=None,\n)\n\nconfig = AutoEasyDelConfig.from_pretrained(\n huggingface_model_repo_id\n)\n\ntokenizer = AutoTokenizer.from_pretrained(\n huggingface_model_repo_id,\n trust_remote_code=True\n)\n\nmax_length = config.max_position_embeddings\n\nconfigs_to_initialize_model_class = {\n 'config': config,\n 'dtype': jnp.bfloat16,\n 'param_dtype': jnp.bfloat16,\n 'input_shape': (8, 8)\n}\n```\n\n`EasyDelState` also has `.load_state()` and `.save_state()` with some other usable options like `.free_opt_state()`\nwhich\nfree optimizer state or `.shard_params()` which shard parameters you can read docs in order to find out more about these\noptions.\n\n### Converting to Huggingface and Pytorch\n\nLet see how you can convert a EasyDelMistral Model to Huggingface Pytorch Mistral Model from a trained State\n\n```python\n\nfrom transformers import MistralForCausalLM\nfrom EasyDel import (\n AutoEasyDelConfig,\n EasyDelState,\n easystate_to_huggingface_model\n)\nimport jax\n\nhuggingface_model_repo_id = \"REPO_ID\"\n\nconfig = AutoEasyDelConfig.from_pretrained(\n huggingface_model_repo_id\n)\nwith jax.default_device(jax.devices(\"cpu\")[0]):\n model = easystate_to_huggingface_model(\n state=EasyDelState.load_state(\n \"PATH_TO_CKPT\"\n ), # You can Pass EasyDelState here\n base_huggingface_module=MistralForCausalLM, # type: ignore\n config=config\n )\n\nmodel = model.half() # it's a huggingface model now\n```\n\n### Other Use Cases\n\n`EasyDelState` have a general use you can use it everywhere in easydel for example for a stand-alone model\n, serve, fine-tuning and many other features, it's up to you to test how creative you are \ud83d\ude07.\n\n## Flash Attention and Splash Attention Are Here \ud83e\udd75\n\nhere's a simple example about how can you use Flash Attention in EasyDeL\n\n```python\n# Config is built in config for every model (EasyDelPretrainedConfig)\nconfig.add_basic_configurations(\n attn_mechanism=\"flash\", # Any supported Attention Mechanism\n block_b=1,\n block_q=512,\n block_k=512,\n block_k_major=512\n)\n```\n\n_Flash Attention works on TPU with ease but for gpu there are still some improvements in process._\n\n> [!TIP]\n> use these partition specs in case of not using custom sharding_axis_names and using sequence sharding with flash\n> flash attention\n> ```python\n>query_partition_spec=PartitionSpec((\"dp\", \"fsdp\"), None, \"sp\", \"tp\"),\n>generation_query_partition_spec=PartitionSpec((\"dp\", \"fsdp\"), None, None, \"tp\"),\n>key_partition_spec=PartitionSpec((\"dp\", \"fsdp\"), None, \"sp\", \"tp\"),\n>value_partition_spec=PartitionSpec((\"dp\", \"fsdp\"), None, \"sp\", \"tp\"),\n>attention_partition_spec=PartitionSpec((\"dp\", \"fsdp\"), None,\"sp\", \"tp\"),\n> ```\n\n## EasyDeLXRapTure for layer tuning and LoRA\n\nin case of using LoRA and applying that on the EasyDeL models there are some other things\nthat you might need to config on your own but a lot of things being handled by EasyDeL so let just jump into an example\nfor LoRA fine-tuning section and use _EasyDeLXRapTure_ in for mistral models with flash attention example\n\n```python\nfrom flax.core import FrozenDict\nfrom EasyDel import (\n TrainArguments,\n CausalLanguageModelTrainer,\n AutoEasyDelModelForCausalLM,\n EasyDelOptimizers,\n EasyDelSchedulers,\n EasyDelGradientCheckPointers,\n EasyDeLXRapTureConfig\n)\nfrom datasets import load_dataset\nimport flax\nfrom jax import numpy as jnp\nfrom transformers import AutoTokenizer\n\nhuggingface_repo_id_or_path = \"mistralai/Mistral-7B-Instruct-v0.1\"\n\nmodel, params = AutoEasyDelModelForCausalLM.from_pretrained(huggingface_repo_id_or_path, )\n\nmax_length = 8196\nmodel_parameters = FrozenDict({\"params\": params})\n\ndtype = jnp.bfloat16\nparam_dtype = jnp.bfloat16 # you can change that if you want \n\ntokenizer = AutoTokenizer.from_pretrained(\n huggingface_repo_id_or_path,\n trust_remote_code=True\n)\n\nmodel.config.add_basic_configurations(\n attn_mechanism=\"flash\", # Using FlashAttention\n block_b=1,\n block_q=1024,\n block_k=1024,\n block_k_major=1024,\n)\n\ntokenizer.pad_token = tokenizer.eos_token\nconfigs_to_initialize_model_class = {\n \"config\": model.config,\n \"dtype\": dtype,\n \"param_dtype\": param_dtype,\n \"input_shape\": (1, 1)\n}\n\nrapture = EasyDeLXRapTureConfig(\n parameters=model_parameters,\n lora_dim=64,\n fully_fine_tune_parameters=[\"embed_tokens\"], # Model layer to be fully fine tuned\n lora_fine_tune_parameters=[\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\"], # LoRA Layer Targets you can pass this to none\n # For only Layer Tuning or transfer learning\n verbose=True\n)\n\ntrain_arguments = TrainArguments(\n model_class=type(model),\n model_name=\"EasyDeL-Lora-Example\",\n num_train_epochs=3,\n configs_to_initialize_model_class=configs_to_initialize_model_class,\n learning_rate=1e-4, # Using higher learning rate is recommended\n learning_rate_end=8e-5,\n optimizer=EasyDelOptimizers.ADAMW, # \"adamw\", \"lion\", \"adafactor\" are supported\n scheduler=EasyDelSchedulers.LINEAR,\n # \"linear\",\"cosine\", \"none\" ,\"warm_up_cosine\" and \"warm_up_linear\" are supported\n weight_decay=0.01,\n total_batch_size=512,\n max_training_steps=None, # None to let trainer Decide\n do_train=True,\n do_eval=False, # it's optional but supported \n backend=\"tpu\", # default backed is set to cpu, so you must define you want to use tpu cpu or gpu\n max_length=max_length, # Note that you have to change this in the model config too\n gradient_checkpointing=EasyDelGradientCheckPointers.NOTHING_SAVEABLE,\n sharding_array=(1, -1, 1, 1), # the way to shard model across gpu,cpu or TPUs using sharding array (1, -1, 1, 1)\n # everything training will be in sequence and model parallel automatic and share data between devices\n remove_ckpt_after_load=True,\n gradient_accumulation_steps=1,\n loss_re_mat=\"\",\n dtype=dtype,\n param_dtype=param_dtype,\n rapture_config=rapture,\n merge_lora_rapture_parameters=True # turning this off is still not supported and not recommended to do so\n # What this does ? this will merge the lora parameters with the original model parameters and the end of training\n)\n\n\ndef ultra_chat_prompting_process(\n data_chunk\n):\n user_part = [\n chunk[\"content\"] for chunk in data_chunk[\"messages\"] if chunk[\"role\"] == \"user\"\n ]\n assistant_part = [\n chunk[\"content\"] for chunk in data_chunk[\"messages\"] if chunk[\"role\"] == \"assistant\"\n ]\n\n prompt = \"\"\n\n for uc, ac in zip(user_part, assistant_part):\n prompt += f\"<|user|>\\n{uc}</s>\\n<|assistant|>\\n{ac}</s>\\n\"\n\n return {\"prompt\": prompt}\n\n\ntokenization_process = lambda data_chunk: tokenizer(\n data_chunk[\"prompt\"],\n add_special_tokens=False,\n max_length=max_length,\n padding=\"max_length\"\n)\n\ndataset = load_dataset(\"HuggingFaceH4/ultrachat_200k\")\ndataset_train = dataset[\"train_gen\"].map(ultra_chat_prompting_process, num_proc=12)\ndataset_train = dataset_train.map(\n tokenization_process,\n num_proc=12,\n remove_columns=dataset_train.column_names\n)\n\n# you can do the same for evaluation process dataset\n\ntrainer = CausalLanguageModelTrainer(\n train_arguments,\n dataset_train,\n checkpoint_path=None\n)\n\noutput = trainer.train() # you should not pass the parameters in Trainer.train anymore when\n# you are using LoRA or transfer Learning\nprint(f\"Hey ! , here's where your model saved {output.checkpoint_path}\")\n```\n\n## Contributing\n\nEasyDeL is an open-source project, and contributions are welcome. If you would like to contribute to EasyDeL, please\nfork the repository, make your changes, and submit a pull request. The team behind EasyDeL will review your changes and\nmerge them if they are suitable.\n\n## License \ud83d\udcdc\n\nEasyDeL is an Fully Open-Source released under the Apache v2 license. Please see the LICENSE file in the root directory\nof this project for\nmore information.\n\n## Contact\n\nIf you have any questions or comments about EasyDeL, you can reach out to me\n\n## Citing EasyDeL \ud83e\udd76\n\nTo cite this repository:\n\n```misc\n@misc{Zare Chavoshi_2023,\n title={EasyDeL, an open-source library, is specifically designed to enhance and streamline the training process of machine learning models. It focuses primarily on Jax/Flax and aims to provide convenient and effective solutions for training Flax/Jax Models on TPU/GPU for both Serving and Training purposes.},\n url={https://github.com/erfanzar/EasyDel},\n journal={EasyDeL Easy and Fast DeepLearning with JAX},\n publisher={Erfan Zare Chavoshi},\n author={Zare Chavoshi, Erfan},\n year={2023}\n} \n```\n",
"bugtrack_url": null,
"license": "Apache-2.0",
"summary": "An open-source library to make training faster and more optimized in Jax/Flax",
"version": "0.0.63",
"project_urls": {
"Documentation": "https://erfanzar.github.io/EasyDeL",
"Homepage": "https://github.com/erfanzar/EasyDeL",
"Issues": "https://github.com/erfanzar/EasyDeL/issues"
},
"split_keywords": [
"jax",
" torch",
" deep learning",
" machine learning",
" flax",
" xla"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "ff5fac2e55c505f0f8c1a992e139958a6cba2e5b51f99bf454a95cecd26ee7b6",
"md5": "30f2082ec6cb5cc1cb1a6d584f936c81",
"sha256": "cdac48f8cd53a89914f823c7aa2df3bcdd130f2d0ca1a08844444d1b4e2bfd4e"
},
"downloads": -1,
"filename": "EasyDeL-0.0.63-py3-none-any.whl",
"has_sig": false,
"md5_digest": "30f2082ec6cb5cc1cb1a6d584f936c81",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.8",
"size": 452996,
"upload_time": "2024-04-27T12:47:41",
"upload_time_iso_8601": "2024-04-27T12:47:41.472650Z",
"url": "https://files.pythonhosted.org/packages/ff/5f/ac2e55c505f0f8c1a992e139958a6cba2e5b51f99bf454a95cecd26ee7b6/EasyDeL-0.0.63-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "7f8687337a43b7fd5a984085e7c0c34683db860c61875fec900f66dc271bea28",
"md5": "d59792deb0ea7eb973d152ee3d8202a2",
"sha256": "6d75874e02070a82988444e27e27850709ee22fcb765a5cfa2c8b7815dedde65"
},
"downloads": -1,
"filename": "easydel-0.0.63.tar.gz",
"has_sig": false,
"md5_digest": "d59792deb0ea7eb973d152ee3d8202a2",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.8",
"size": 385258,
"upload_time": "2024-04-27T12:48:14",
"upload_time_iso_8601": "2024-04-27T12:48:14.470855Z",
"url": "https://files.pythonhosted.org/packages/7f/86/87337a43b7fd5a984085e7c0c34683db860c61875fec900f66dc271bea28/easydel-0.0.63.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-04-27 12:48:14",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "erfanzar",
"github_project": "EasyDeL",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [
{
"name": "chex",
"specs": [
[
">=",
"0.1.7"
]
]
},
{
"name": "typing",
"specs": [
[
"~=",
"3.7.4.3"
]
]
},
{
"name": "jax",
"specs": [
[
">=",
"0.4.20"
]
]
},
{
"name": "jaxlib",
"specs": [
[
">=",
"0.4.20"
]
]
},
{
"name": "flax",
"specs": [
[
">=",
"0.7.5"
]
]
},
{
"name": "fjformer",
"specs": [
[
">=",
"0.0.50"
]
]
},
{
"name": "transformers",
"specs": [
[
">=",
"4.34.0"
]
]
},
{
"name": "einops",
"specs": [
[
"~=",
"0.6.1"
]
]
},
{
"name": "optax",
"specs": [
[
"~=",
"0.1.7"
]
]
},
{
"name": "msgpack",
"specs": [
[
"~=",
"1.0.7"
]
]
},
{
"name": "ipython",
"specs": [
[
"~=",
"8.17.2"
]
]
},
{
"name": "tqdm",
"specs": [
[
"~=",
"4.64.1"
]
]
},
{
"name": "pydantic",
"specs": [
[
"==",
"2.5.3"
]
]
},
{
"name": "datasets",
"specs": [
[
"~=",
"2.14.6"
]
]
},
{
"name": "setuptools",
"specs": [
[
"~=",
"68.1.2"
]
]
},
{
"name": "gradio",
"specs": [
[
"~=",
"4.25.0"
]
]
},
{
"name": "termcolor",
"specs": []
},
{
"name": "wandb",
"specs": [
[
">=",
"0.15.9"
]
]
},
{
"name": "tensorboard",
"specs": []
},
{
"name": "torch",
"specs": [
[
">=",
"2.1.0"
]
]
},
{
"name": "numpy",
"specs": [
[
"~=",
"1.26.2"
]
]
},
{
"name": "uvicorn",
"specs": [
[
"~=",
"0.23.2"
]
]
},
{
"name": "fastapi",
"specs": [
[
"~=",
"0.104.1"
]
]
},
{
"name": "pydantic-core",
"specs": [
[
"==",
"2.14.6"
]
]
},
{
"name": "tensorflow_datasets",
"specs": [
[
"<=",
"4.9.0"
]
]
},
{
"name": "tensorflow",
"specs": [
[
"~=",
"2.16.1"
]
]
},
{
"name": "scipy",
"specs": [
[
"==",
"1.10.1"
]
]
}
],
"lcname": "easydel"
}
JSON
