Skip to content

Commit

Permalink
fixes 6
Browse files Browse the repository at this point in the history
  • Loading branch information
Almaz Dautov committed Nov 12, 2024
1 parent fce6c51 commit ca5a39e
Show file tree
Hide file tree
Showing 10 changed files with 2,186 additions and 1,609 deletions.
3,495 changes: 1,911 additions & 1,584 deletions poetry.lock
View file

Large diffs are not rendered by default.

4 changes: 2 additions & 2 deletions pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ deepspeed = "^0.15.3"
accelerate = "^1.1.1"
vllm = "^0.6.3.post1"
allenai-common = "^1.1.2"

liger-kernel-nightly = {version = "^0.4.0.dev20241109021846", optional = true}

[tool.poetry.dev-dependencies]
coverage = "^7.2.7"
Expand All @@ -71,7 +71,7 @@ annotated-types = "^0.6.0"

[tool.poetry.extras]
deepspeed = ["deepspeed", "accelerate"]
gpu = ["faiss-gpu", "xformers", "vllm"]
gpu = ["faiss-gpu", "xformers", "vllm", "liger-kernel-nightly"]

[build-system]
requires = ["poetry-core>=1.0.0"]
Expand Down
4 changes: 2 additions & 2 deletions tests/fixtures/configs/train/reinforce/base.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@
"trainer_settings": {
"actor_settings": {
"actor_type": "distributed_vllm",
"vllm_num_engines": 1,
"vllm_num_engines": 6,
"vllm_tensor_parallel_size": 1
},
"deepspeed": "tests/fixtures/configs/train/reinforce/deepspeed_cfg.json",
Expand Down Expand Up @@ -151,7 +151,7 @@
"max_grad_norm": 2,
"save_total_limit": 11,
"dataloader_num_workers": 12,
"stop_token": "<|eot_id|>",
"stop_token": "<end_of_turn>",
"temperature": 0.7,
"non_eos_penalty": true,
"penalty_reward_value": -1,
Expand Down
2 changes: 1 addition & 1 deletion tests/fixtures/configs/train/reinforce/deepspeed_cfg.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
},

"zero_optimization": {
"stage": 2,
"stage": 3,
"offload_param": {
"device": "none"
},
Expand Down
195 changes: 195 additions & 0 deletions tests/fixtures/configs/train/reinforce/gemma.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,195 @@
{
"train_dataset_settings": {
"sources": [
{
"name": "train_chat",
"records_path": "tests/fixtures/datasets/chat/train_chat.jsonl",
"sample_rate": 1.0
}
],
"prompt_template": {
"role_tag_mapping": {
"bot": "assistant",
"user": "user",
"system": "system"
},
"prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
"suffix_template": "<|eot_id|>"
},
"dataset_type": "chat",
"max_tokens_count": 2000,
"only_answer_loss": true
},
"val_dataset_settings": {
"sources": [
{
"name": "val_chat",
"records_path": "tests/fixtures/datasets/chat/train_chat.jsonl",
"sample_rate": 1.0
}
],
"prompt_template": {
"role_tag_mapping": {
"bot": "assistant",
"user": "user",
"system": "system"
},
"prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
"suffix_template": "<|eot_id|>"
},
"dataset_type": "chat",
"max_tokens_count": 2000,
"only_answer_loss": true
},
"cherry_pick_settings": {
"generator_transformers_settings": {
"num_beams": 1,
"do_sample": false,
"stop_strings": "</RS>",
"max_new_tokens": 8
},
"custom_generation_settings": {
"skip_special_tokens": false
},
"dataset_settings": {
"sources": [
{
"name": "chat_test",
"records_path": "tests/fixtures/datasets/chat/train_chat.jsonl",
"num_samples": 2
}
],
"prompt_template": {
"role_tag_mapping": {
"bot": "assistant",
"user": "user",
"system": "system"
},
"prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
"suffix_template": "<|eot_id|>"
},
"dataset_type": "chat",
"max_tokens_count": 150,
"only_answer_loss": true
},
"metric_settings": [
{
"type": "length",
"parameters": {
"need_average": [
true
]
}
},
{
"type": "kl",
"parameters": {
"need_average": [
true
],
"ref_logits_type": "sft"
}
},
{
"type": "kl",
"parameters": {
"need_average": [
true
],
"ref_logits_type": "reference"
}
}
]
},
"reward_model_settings": {
"model_path": "/from_s3/llama-3.2-1B",
"model_type": "seq_cls",
"model_kwargs": {
"num_labels": 1,
"attn_implementation": "flash_attention_2"
},
"transformers_settings": {}
},
"model_settings": {
"model_path": "/from_s3/llama-3.2-1B",
"model_type": "causal",
"liger_kernels_settings": {
"use_rope": true,
"use_cross_entropy": false,
"use_geglu": true,
"use_fused_linear_cross_entropy": false,
"use_rms_norm": false
},
"model_kwargs": {
"attn_implementation": "flash_attention_2"
},
"transformers_settings": {}
},
"tokenizer_settings": {
"tokenizer_kwargs": {
"padding_side": "left"
}
},
"trainer_settings": {
"actor_settings": {
"actor_type": "distributed_vllm",
"vllm_num_engines": 6,
"vllm_tensor_parallel_size": 1
},
"deepspeed": "tests/fixtures/configs/train/reinforce/deepspeed_cfg.json",
"gradient_checkpointing": true,
"gradient_checkpointing_kwargs": {
"use_reentrant": false
},
"critic_type": "ray_transformers",
"reward_processor_type": "rloo",
"evaluation_strategy": "steps",
"num_generations": 4,
"per_device_train_batch_size": 8,
"per_device_eval_batch_size": 8,
"gradient_accumulation_steps": 10,
"adam_beta1": 0.9,
"adam_beta2": 0.95,
"adam_epsilon": 0.00001,
"eval_steps": 50,
"save_steps": 100,
"save_strategy": "steps",
"load_best_model_at_end": false,
"logging_steps": 1,
"learning_rate": 0.00001,
"max_steps": 1001,
"lr_scheduler_type": "constant_with_warmup",
"warmup_ratio": 0.03,
"fp16": false,
"bf16": true,
"optim": "adamw_torch",
"weight_decay": 0.0,
"max_grad_norm": 2,
"save_total_limit": 11,
"dataloader_num_workers": 12,
"stop_token": "<end_of_turn>",
"temperature": 0.7,
"non_eos_penalty": true,
"penalty_reward_value": -1,
"clip_rewards_min": -1e+8,
"clip_rewards_max": 1e+8,
"whiten_rewards": false,
"kl_coef": 0.05,
"mean_baseline_coef": 0.95,
"num_samples_for_reward_stats": 1000,
"no_cuda": false
},
"special_tokens_settings": {
"bos_token": "<s>",
"eos_token": "</s>",
"pad_token": "<unk>"
},
"logging_settings": {
"project_name": "alignment",
"run_name": "sft",
"entity": "turbo-alignment"
},
"seed": 0,
"log_path": "train_output"
}

55 changes: 55 additions & 0 deletions tests/fixtures/configs/train/reinforce/stage3.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
{
"fp16": {
"enabled": false
},
"bf16": {
"enabled": true
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupDecayLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto",
"total_num_steps": "auto"
}
},

"zero_optimization": {
"stage": 3,
"offload_param": {
"device": "none"
},
"offload_optimizer": {
"device": "none",
"pin_memory": true
},
"allgather_partitions": true,
"allgather_bucket_size": 1e5,
"overlap_comm": false,
"reduce_scatter": true,
"reduce_bucket_size": "auto",
"contiguous_gradients": true,
"sub_group_size": 1e5,
"stage3_param_persistence_threshold": "auto",
"stage3_prefetch_bucket_size": 0,
"stage3_max_live_parameters": 0,
"stage3_max_reuse_distance": 0,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
2 changes: 2 additions & 0 deletions turbo_alignment/settings/model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ class LigerKernelSettings(ExtraFieldsNotAllowedBaseModel):
use_rope: bool = True
use_cross_entropy: bool = True
use_geglu: bool = True
use_fused_linear_cross_entropy: bool = False
use_rms_norm: bool = False


class PreTrainedModelSettings(ExtraFieldsNotAllowedBaseModel):
Expand Down
2 changes: 1 addition & 1 deletion turbo_alignment/settings/tf/trainer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from turbo_alignment.settings.base import ExtraFieldsNotAllowedBaseModel


class TrainerSettings(ExtraFieldsNotAllowedBaseModel):
class TrainerSettings(ExtraFieldsNotAllowedBaseModel):
evaluation_strategy: str = 'steps'
save_strategy: str = 'steps'
per_device_train_batch_size: int = 4
Expand Down
Loading

0 comments on commit ca5a39e

Please sign in to comment.