diff --git a/.gitignore b/.gitignore index c243024..8c28ce3 100644 --- a/.gitignore +++ b/.gitignore @@ -175,6 +175,7 @@ debug.py wandb/ nohup.out lm-evaluation-harness/ +bigcode-evaluation-harness/ results/**/*.json results/**/*.jsonl results/**/*.db diff --git a/README.md b/README.md index 8813a32..b276a78 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,11 @@ bash scripts/data.sh git clone https://github.com/EleutherAI/lm-evaluation-harness.git cd lm-evaluation-harness pip install -e . +# commit: 9cfa52b +git clone https://github.com/bigcode-project/bigcode-evaluation-harness.git +cd bigcode-evaluation-harness +# change `pyext==0.5` in `bigcode-evaluation-harness/requirements.txt`, ref: https://github.com/bigcode-project/bigcode-evaluation-harness/pull/181 +pip install -e . ``` ## 📃 TODO diff --git a/scripts/eval.sh b/scripts/eval.sh deleted file mode 100644 index 4f41b37..0000000 --- a/scripts/eval.sh +++ /dev/null @@ -1,96 +0,0 @@ -# nohup srun -p MoE --gres gpu:1 bash scripts/eval.sh all /mnt/petrelfs/share_data/quxiaoye/models/Sheared-LLaMA-2.7B True results/Sheared-LLaMA-2.7B 1>logs/eval-all-Sheared-LLaMA-2.7B.log 2>&1 & - -mmlu() { - # MMLU: https://github.com/princeton-nlp/LLM-Shearing/blob/20ebd2645a8ff5fa65874e1347f9891b80e01805/icl_eval/run_eval.sh#L18 - MODEL=$1 - TRUST_REMOTE_CODE=$2 - RESULT_DIR=$3 - mkdir -p $RESULT_DIR - - lm_eval \ - --model hf \ - --model_args pretrained=$MODEL,trust_remote_code=$TRUST_REMOTE_CODE \ - --tasks mmlu_computer_security,mmlu_high_school_chemistry,mmlu_philosophy,mmlu_elementary_mathematics,mmlu_prehistory,mmlu_formal_logic,mmlu_high_school_mathematics,mmlu_econometrics,mmlu_moral_scenarios,mmlu_college_mathematics,mmlu_high_school_government_and_politics,mmlu_us_foreign_policy,mmlu_high_school_world_history,mmlu_conceptual_physics,mmlu_college_medicine,mmlu_international_law,mmlu_abstract_algebra,mmlu_logical_fallacies,mmlu_machine_learning,mmlu_medical_genetics,mmlu_public_relations,mmlu_college_biology,mmlu_marketing,mmlu_electrical_engineering,mmlu_anatomy,mmlu_high_school_us_history,mmlu_high_school_biology,mmlu_miscellaneous,mmlu_high_school_psychology,mmlu_sociology,mmlu_business_ethics,mmlu_high_school_geography,mmlu_human_aging,mmlu_high_school_statistics,mmlu_moral_disputes,mmlu_professional_psychology,mmlu_global_facts,mmlu_college_physics,mmlu_nutrition,mmlu_high_school_macroeconomics,mmlu_world_religions,mmlu_professional_medicine,mmlu_high_school_computer_science,mmlu_college_chemistry,mmlu_human_sexuality,mmlu_high_school_microeconomics,mmlu_astronomy,mmlu_professional_accounting,mmlu_high_school_european_history,mmlu_jurisprudence,mmlu_professional_law,mmlu_high_school_physics,mmlu_virology,mmlu_management,mmlu_college_computer_science,mmlu_clinical_knowledge,mmlu_security_studies \ - --num_fewshot 5 \ - --device cuda:0 \ - --batch_size auto \ - --verbosity DEBUG \ - --output_path $RESULT_DIR/mmlu.json -} - -bbh() { - # Big Bench Hard (BBH): https://arxiv.org/pdf/2210.09261.pdf - MODEL=$1 - TRUST_REMOTE_CODE=$2 - RESULT_DIR=$3 - mkdir -p $RESULT_DIR - - lm_eval \ - --log_samples \ - --model hf \ - --model_args pretrained=$MODEL,trust_remote_code=$TRUST_REMOTE_CODE \ - --tasks bbh_fewshot_boolean_expressions,bbh_fewshot_causal_judgement,bbh_fewshot_date_understanding,bbh_fewshot_disambiguation_qa,bbh_fewshot_dyck_languages,bbh_fewshot_formal_fallacies,bbh_fewshot_geometric_shapes,bbh_fewshot_hyperbaton,bbh_fewshot_logical_deduction_five_objects,bbh_fewshot_logical_deduction_seven_objects,bbh_fewshot_logical_deduction_three_objects,bbh_fewshot_movie_recommendation,bbh_fewshot_multistep_arithmetic_two,bbh_fewshot_navigate,bbh_fewshot_object_counting,bbh_fewshot_penguins_in_a_table,bbh_fewshot_reasoning_about_colored_objects,bbh_fewshot_ruin_names,bbh_fewshot_salient_translation_error_detection,bbh_fewshot_snarks,bbh_fewshot_sports_understanding,bbh_fewshot_temporal_sequences,bbh_fewshot_tracking_shuffled_objects_five_objects,bbh_fewshot_tracking_shuffled_objects_seven_objects,bbh_fewshot_tracking_shuffled_objects_three_objects,bbh_fewshot_web_of_lies,bbh_fewshot_word_sorting \ - --device cuda:0 \ - --batch_size auto \ - --verbosity DEBUG \ - --output_path $RESULT_DIR/bbh.json -} - -reasoning() { - MODEL=$1 - TRUST_REMOTE_CODE=$2 - RESULT_DIR=$3 - mkdir -p $RESULT_DIR - - lm_eval \ - --log_samples \ - --model hf \ - --model_args pretrained=$MODEL,trust_remote_code=$TRUST_REMOTE_CODE \ - --tasks gsm8k_cot \ - --device cuda:0 \ - --batch_size auto \ - --verbosity DEBUG \ - --output_path $RESULT_DIR/reasoning.json -} - -qa() { - MODEL=$1 - TRUST_REMOTE_CODE=$2 - RESULT_DIR=$3 - mkdir -p $RESULT_DIR - - lm_eval \ - --log_samples \ - --model hf \ - --model_args pretrained=$MODEL,trust_remote_code=$TRUST_REMOTE_CODE \ - --tasks arc_easy,arc_challenge,boolq \ - --num_fewshot 0 \ - --device cuda:0 \ - --batch_size auto \ - --verbosity DEBUG \ - --output_path $RESULT_DIR/qa.json -} - -EVAL_TASK=$1 -shift 1 -start=$(date +%s) -case $EVAL_TASK in - mmlu) - mmlu $* ;; - bbh) - bbh $* ;; - reasoning) - reasoning $* ;; - qa) - qa $* ;; - all) - mmlu $* - bbh $* - reasoning $* - qa $* - ;; - *) - echo "$EVAL_TASK not recognized!";; -esac -end=$(date +%s) -echo "Elapsed Time: $(($end-$start)) seconds" diff --git a/scripts/four_mix/freeze_gate.sh b/scripts/four_mix/freeze_gate.sh index d94d78c..70afb8e 100644 --- a/scripts/four_mix/freeze_gate.sh +++ b/scripts/four_mix/freeze_gate.sh @@ -83,8 +83,11 @@ num_gpus=4 python -m src.eval.gen_mt_ans \ --model-path $output_dir \ - --model-id $task_name \ - --num-gpus-total $num_gpus + --model-id $task_name + + python -m src.eval.gen_alpaca_eval_ans \ + --model-path $output_dir \ + --model-id $task_name } # nohup srun -p MoE --ntasks-per-node=1 --cpus-per-task=16 --mem=128G --nodes=1 --gres=gpu:4 bash "/mnt/petrelfs/zhutong/adaptive-sft-for-moe/scripts/one_data_steps_dynamic.sh" "llama_moe_orca_epochs_cluster_4" "auto" "/mnt/petrelfs/zhutong/llama-moe-models/LLaMA-MoE-v1-3_5B-2_8-new" "data/open_orca_clustered/4" "data/open_orca_clustered_eval/4" 1>logs/llama_moe_orca_cluster_4_dynamic.log 2>&1 & diff --git a/scripts/gen_mt_bench_ans.sh b/scripts/gen_mt_bench_ans.sh deleted file mode 100644 index f251644..0000000 --- a/scripts/gen_mt_bench_ans.sh +++ /dev/null @@ -1,32 +0,0 @@ -#!/usr/bin/bash - -#SBATCH --job-name=moe_gen -#SBATCH --output=logs/%x-%j.log -#SBATCH --error=logs/%x-%j.log - -#SBATCH --partition=MoE -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=16 -#SBATCH --mem=64G - -#SBATCH --nodes=1 -#SBATCH --gres=gpu:1 -#SBATCH --quotatype=auto - -{ - # python -m fastchat.llm_judge.gen_model_answer \ - # --model-path outputs/sheared_llama_sharegpt/moe_sft-2411306 \ - # --model-id sheared_llama_sharegpt - - # python -m fastchat.llm_judge.gen_model_answer \ - # --model-path outputs/sheared_llama_uniform_mix/moe_sft-2421072 \ - # --model-id sheared_llama_uniform_mix - - bash scripts/cp_model_files.sh outputs/llama_moe/moe_sft-2409782 - python -m fastchat.llm_judge.gen_model_answer \ - --model-path outputs/llama_moe/moe_sft-2409782 \ - --model-id llama_moe_uniform_mix -} - -# nohup srun -p MoE -n1 -N1 --gres=gpu:1 --quotatype spot python -m fastchat.llm_judge.gen_model_answer --model-path outputs/sheared_llama_sharegpt/moe_sft-2411306 --model-id sheared_llama_sharegpt 1>logs/mt_bench_gen_sheared_llama_sharegpt.log 2>&1 & -# nohup srun -p MoE -n1 -N1 --gres=gpu:1 --quotatype spot python -m fastchat.llm_judge.gen_model_answer --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/llama_moe_sharegpt/moe_sft-2411309 --model-id llama_moe_sharegpt 1>logs/mt_bench_gen_llama_moe_sharegpt.log 2>&1 & diff --git a/scripts/multi.sh b/scripts/multi.sh index bcd83b8..e399761 100644 --- a/scripts/multi.sh +++ b/scripts/multi.sh @@ -100,5 +100,8 @@ nohup srun -p MoE --ntasks-per-node=1 --cpus-per-task=16 --mem=128G --nodes=1 -- nohup srun -p MoE --gres gpu:1 python -m src.eval.gen_mt_ans --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048/llama_moe_four_mix_uniform/bash-2485396 --model-id llama_moe_four_mix_uniform 1>logs/gen_mt_ans-llama_moe_four_mix_uniform.log 2>&1 & nohup srun -p MoE --gres gpu:1 python -m src.eval.gen_mt_ans --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048/sheared_four_mix_uniform/bash-2485397 --model-id sheared_four_mix_uniform 1>logs/gen_mt_ans-sheared_four_mix_uniform.log 2>&1 & -nohup srun -p MoE --gres gpu:1 python -m src.eval.get_alpaca_eval_ans --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048/llama_moe_four_mix_uniform/bash-2485396 --model-id llama_moe_four_mix_uniform 1>logs/gen_alpaca_eval-llama_moe_four_mix_uniform.log 2>&1 & -nohup srun -p MoE --gres gpu:1 python -m src.eval.get_alpaca_eval_ans --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048/sheared_four_mix_uniform/bash-2485397 --model-id sheared_four_mix_uniform 1>logs/gen_alpaca_eval-sheared_four_mix_uniform.log 2>&1 & +nohup srun -p MoE --gres gpu:1 python -m src.eval.gen_alpaca_eval_ans --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048/llama_moe_four_mix_uniform/bash-2485396 --model-id llama_moe_four_mix_uniform 1>logs/gen_alpaca_eval-llama_moe_four_mix_uniform.log 2>&1 & +nohup srun -p MoE --gres gpu:1 python -m src.eval.gen_alpaca_eval_ans --model-path /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048/sheared_four_mix_uniform/bash-2485397 --model-id sheared_four_mix_uniform 1>logs/gen_alpaca_eval-sheared_four_mix_uniform.log 2>&1 & + +nohup srun -p MoE --gres gpu:1 bash scripts/eval/eval.sh reasoning /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048_dynamic_remove_padding_tokens/llama_moe_four_mix_wo_pad_wo_gate_noise/moe_sft-2492650 True results/llama_moe_four_mix_wo_pad_wo_gate_noise 1>logs/eval-reasoning-llama_moe_four_mix_wo_pad_wo_gate_noise.log 2>&1 & +nohup srun -p MoE --gres gpu:1 bash scripts/eval/eval.sh reasoning /mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048_dynamic_remove_padding_tokens/llama_moe_four_mix_wo_pad/moe_sft-2491633 True results/llama_moe_four_mix_wo_pad 1>logs/eval-reasoning-llama_moe_four_mix_wo_pad.log 2>&1 & diff --git a/src/callbacks.py b/src/callbacks.py index a750f69..e9d0c04 100644 --- a/src/callbacks.py +++ b/src/callbacks.py @@ -6,6 +6,7 @@ import torch import numpy as np from loguru import logger from transformers.trainer_callback import TrainerCallback, TrainerState, TrainerControl +from transformers.utils import is_flash_attn_2_available from src.utils.config import TrainingArguments from src.utils.io import append_jsonlines @@ -22,6 +23,7 @@ class AdaptiveSamplingCallback(TrainerCallback): criterion: Optional[Literal["min", "max", "mean"]] = "mean", sim_type: Optional[Literal["cos", "l2"]] = "cos", ): + assert is_flash_attn_2_available(), "Make sure you have flash-attn installed" self.criterion = criterion self.sim_type = sim_type self.prob_map = {} @@ -74,8 +76,8 @@ class AdaptiveSamplingCallback(TrainerCallback): cls, ori_weights: np.ndarray, delta: np.ndarray, - eta: float = 1.0, - c: float = 1e-4, + eta: float = 10.0, + c: float = 5e-2, ) -> np.ndarray: def _softmax(vec: np.ndarray) -> np.ndarray: exps = np.exp(vec - np.max(vec)) diff --git a/src/core/train.py b/src/core/train.py index 2be5558..9b1f694 100644 --- a/src/core/train.py +++ b/src/core/train.py @@ -7,13 +7,12 @@ from loguru import logger from src.utils.config import ModelArguments, DataArguments, TrainingArguments from src.data import ( SubDirWeightedPackedJsonlDataset, - get_uniform_sampling_ratio, fault_tolerance_data_collator, CachedJsonlDataset, get_cached_datasets_from_dir, ) from src.utils.io import trainer_save_model_safe -from src.models import LlamaMoEForCausalLM, LlamaMoEConfig +from src.models import LlamaMoEForCausalLM, LlamaMoEConfig, DeepseekConfig, DeepseekForCausalLM from src.trainer import GateLoadRecordingTrainer from src.callbacks import AdaptiveSamplingCallback @@ -36,6 +35,9 @@ def get_model_and_tokenizer( elif model_type == "llama_moe": ConfigClass = LlamaMoEConfig ModelClass = LlamaMoEForCausalLM + elif model_type == "deepseek": + ConfigClass = DeepseekConfig + ModelClass = DeepseekForCausalLM else: raise ValueError(f"Unknown model type: {model_type}") @@ -54,6 +56,21 @@ def get_model_and_tokenizer( config.update(additional_config) logger.info("Config ready") + tokenizer = transformers.AutoTokenizer.from_pretrained( + model_name_or_path, + cache_dir=cache_dir, + model_max_length=model_max_length, + padding_side=padding_side, + use_fast=False, + trust_remote_code=trust_remote_code, + ) + if tokenizer.pad_token is None: + if tokenizer.unk_token is not None: + tokenizer.pad_token = tokenizer.unk_token + else: + tokenizer.pad_token = tokenizer.eos_token + logger.info(f"tokenizer ready, pad_token: {tokenizer.pad_token}") + # Load model and tokenizer model = ModelClass.from_pretrained( model_name_or_path, @@ -65,18 +82,6 @@ def get_model_and_tokenizer( ) logger.info("model ready") - tokenizer = transformers.AutoTokenizer.from_pretrained( - model_name_or_path, - cache_dir=cache_dir, - model_max_length=model_max_length, - padding_side=padding_side, - use_fast=False, - trust_remote_code=trust_remote_code, - ) - if tokenizer.pad_token != tokenizer.unk_token: - tokenizer.pad_token = tokenizer.unk_token - logger.info("tokenizer ready") - return model, tokenizer @@ -117,7 +122,9 @@ def train(): train_dataset = SubDirWeightedPackedJsonlDataset( data_args.dataset_dir_or_path, tokenizer, - prob_map=get_uniform_sampling_ratio(data_args.dataset_dir_or_path), + # prob_map=get_uniform_sampling_ratio(data_args.dataset_dir_or_path), + # prob_map={"code": 0.25119094959816823, "math": 0.2674581878910902, "orca": 0.243050776175138, "sharegpt": 0.23830008633560357}, + prob_map=data_args.prob_map, seed=training_args.seed, ) elif datapath.is_file(): diff --git a/src/data.py b/src/data.py index d783a21..a1a8ff7 100644 --- a/src/data.py +++ b/src/data.py @@ -20,6 +20,7 @@ def preprocess( instances, tokenizer: transformers.PreTrainedTokenizer, ) -> Dict: + tokenizer_legacy = getattr(tokenizer, "legacy", None) conv = Conversation() roles = {"human": conv.roles[0], "gpt": conv.roles[1]} @@ -72,7 +73,7 @@ def preprocess( # "-2" is hardcoded for the Llama tokenizer to make the offset correct. instruction_len = len(tokenizer(parts[0]).input_ids) - 2 - if i != 0 and not tokenizer.legacy: + if i != 0 and not tokenizer_legacy: # The legacy and non-legacy modes handle special tokens differently instruction_len -= 1 @@ -80,7 +81,7 @@ def preprocess( target[cur_len : cur_len + instruction_len] = IGNORE_TOKEN_ID cur_len += turn_len - if i != 0 and not tokenizer.legacy: + if i != 0 and not tokenizer_legacy: # The legacy and non-legacy modes handle special tokens differently cur_len -= 1 diff --git a/src/eval/get_alpaca_eval_ans.py b/src/eval/get_alpaca_eval_ans.py deleted file mode 100644 index 1ff3e5e..0000000 --- a/src/eval/get_alpaca_eval_ans.py +++ /dev/null @@ -1,113 +0,0 @@ -import argparse -from pathlib import Path - -import torch -import datasets -from tqdm import tqdm - -from src.core.train import get_model_and_tokenizer -from src.utils.conversation import Conversation -from src.utils.io import dump_json - - -@torch.inference_mode() -def run_eval(model_path, model_id, max_new_tokens): - model, tokenizer = get_model_and_tokenizer( - "auto", - model_path, - torch_dtype=torch.bfloat16, - trust_remote_code=True, - ) - model.cuda() - model.eval() - - conv = Conversation() - outputs = [] - eval_set = datasets.load_dataset("tatsu-lab/alpaca_eval", "alpaca_eval")["eval"] - for example in tqdm(eval_set, desc="Eval"): - conv.append_message(conv.roles[0], example["instruction"]) - conv.append_message(conv.roles[1], None) - prompt = conv.get_prompt() - input_ids = tokenizer([prompt], return_tensors="pt").input_ids - conv.clear_msg() - # generate here is a placeholder for your models generations - output_ids = model.generate( - input_ids.cuda(), - do_sample=False, - temperature=0.0, - max_new_tokens=max_new_tokens, - ) - if model.config.is_encoder_decoder: - output_ids = output_ids[0] - else: - output_ids = output_ids[0][len(input_ids[0]) :] # noqa: E203 - # be consistent with the template's stop_token_ids - if conv.stop_token_ids: - stop_token_ids_index = [ - i - for i, id in enumerate(output_ids) - if id in conv.stop_token_ids - ] - if len(stop_token_ids_index) > 0: - output_ids = output_ids[: stop_token_ids_index[0]] - - output = tokenizer.decode( - output_ids, - spaces_between_special_tokens=False, - ) - if conv.stop_str and isinstance(conv.stop_str, list): - stop_str_indices = sorted( - [ - output.find(stop_str) - for stop_str in conv.stop_str - if output.find(stop_str) > 0 - ] - ) - if len(stop_str_indices) > 0: - output = output[: stop_str_indices[0]] - elif conv.stop_str and output.find(conv.stop_str) > 0: - output = output[: output.find(conv.stop_str)] - - for special_token in tokenizer.special_tokens_map.values(): - if isinstance(special_token, list): - for special_tok in special_token: - output = output.replace(special_tok, "") - else: - output = output.replace(special_token, "") - output = output.strip() - - if conv.name == "xgen" and output.startswith("Assistant:"): - output = output.replace("Assistant:", "", 1).strip() - - example["output"] = output - outputs.append(example) - - outpath = Path("results/alpaca_eval") / f"{model_id}.json" - dump_json(outputs, outpath, indent=2) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--model-path", - type=str, - required=True, - help="The path to the weights. This can be a local folder or a Hugging Face repo ID.", - ) - parser.add_argument( - "--model-id", type=str, required=True, help="A custom name for the model." - ) - parser.add_argument( - "--max-new-token", - type=int, - default=1024, - help="The maximum number of new generated tokens.", - ) - - args = parser.parse_args() - - run_eval( - model_path=args.model_path, - model_id=args.model_id, - max_new_tokens=args.max_new_token, - ) diff --git a/src/eval/show.py b/src/eval/show.py index d500054..ea0c210 100644 --- a/src/eval/show.py +++ b/src/eval/show.py @@ -55,13 +55,13 @@ def collect_results(result_dir: str, verbose: bool = True) -> dict: avg = sum(vals) / len(vals) tot_vals.append(avg) if verbose: - logger.info(f"task: {name}, num: {len(tasks.split(','))}, avg: {avg:.3%}") + logger.info(f"task: {name}, num: {len(tasks.split(','))}, avg: {100 * avg:.3f} %") if len(tot_vals) == 0: tot_avg = 0.0 else: tot_avg = sum(tot_vals) / len(tot_vals) - logger.info(f"total avg: {tot_avg:.3%}") + logger.info(f"total avg: {100 * tot_avg:.3f} %") if __name__ == "__main__": diff --git a/src/models/deepseek/modeling_deepseek.py b/src/models/deepseek/modeling_deepseek.py index 1dae56e..20498b2 100644 --- a/src/models/deepseek/modeling_deepseek.py +++ b/src/models/deepseek/modeling_deepseek.py @@ -20,6 +20,7 @@ """ PyTorch DeepSeek model.""" import math import warnings +from dataclasses import dataclass from typing import List, Optional, Tuple, Union import torch @@ -297,7 +298,7 @@ class DeepseekMLP(nn.Module): self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] - def forward(self, x): + def forward(self, x, **kwargs): if self.config.pretraining_tp > 1: slice = self.intermediate_size // self.config.pretraining_tp gate_proj_slices = self.gate_proj.weight.split(slice, dim=0) @@ -328,7 +329,9 @@ class DeepseekMLP(nn.Module): else: down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj + bsz, seq_len, _ = x.shape + load = torch.zeros(bsz * seq_len, self.config.n_routed_experts) + return down_proj, load class MoEGate(nn.Module): @@ -356,7 +359,10 @@ class MoEGate(nn.Module): init.kaiming_uniform_(self.weight, a=math.sqrt(5)) def forward(self, hidden_states): - bsz, seq_len, h = hidden_states.shape + if len(hidden_states.shape) == 2: + bsz, h = hidden_states.shape + else: + bsz, seq_len, h = hidden_states.shape ### compute gating score hidden_states = hidden_states.view(-1, h) logits = F.linear(hidden_states, self.weight, None) @@ -404,7 +410,10 @@ class MoEGate(nn.Module): aux_loss = (Pi * fi).sum() * self.alpha else: aux_loss = None - return topk_idx, topk_weight, aux_loss + _zeros = torch.zeros_like(logits) + _scores_filtered = _zeros.scatter(dim=1, index=topk_idx, src=topk_weight) + load = (_scores_filtered > 0).sum(0) + return topk_idx, topk_weight, aux_loss, load class AddAuxiliaryLoss(torch.autograd.Function): @@ -450,10 +459,19 @@ class DeepseekMoE(nn.Module): config=config, intermediate_size=intermediate_size ) - def forward(self, hidden_states): + def forward(self, hidden_states, attention_mask=None): + bsz, seq_len, hsz = hidden_states.shape + hidden_states = hidden_states.reshape(-1, hsz) + flattened_mask = None + flattened_shape = None + if attention_mask is not None and len(attention_mask.shape) == 2: + flattened_mask = attention_mask.flatten() + flattened_shape = flattened_mask.shape + hidden_states = hidden_states[flattened_mask.bool()] + identity = hidden_states orig_shape = hidden_states.shape - topk_idx, topk_weight, aux_loss = self.gate(hidden_states) + topk_idx, topk_weight, aux_loss, load = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) flat_topk_idx = topk_idx.view(-1) if self.training: @@ -472,7 +490,15 @@ class DeepseekMoE(nn.Module): ).view(*orig_shape) if self.config.n_shared_experts is not None: y = y + self.shared_experts(identity) - return y + + if flattened_mask is not None: + _y = torch.zeros(flattened_shape + (hsz,), dtype=y.dtype, device=y.device) + _y[flattened_mask.bool()] = y + y = _y + + y = y.reshape(bsz, seq_len, hsz) + + return y, load @torch.no_grad() def moe_infer(self, x, flat_expert_indices, flat_expert_weights): @@ -1163,7 +1189,7 @@ class DeepseekDecoderLayer(nn.Module): # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states, load = self.mlp(hidden_states, attention_mask=attention_mask) hidden_states = residual + hidden_states outputs = (hidden_states,) @@ -1174,6 +1200,8 @@ class DeepseekDecoderLayer(nn.Module): if use_cache: outputs += (present_key_value,) + outputs += (load,) + return outputs @@ -1220,6 +1248,11 @@ class DeepseekPreTrainedModel(PreTrainedModel): module.weight.data[module.padding_idx].zero_() +@dataclass +class BaseMoEModelOutputWithPast(BaseModelOutputWithPast): + gate_load: Optional[torch.Tensor] = None + + Deepseek_INPUTS_DOCSTRING = r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1429,6 +1462,7 @@ class DeepseekModel(DeepseekPreTrainedModel): # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None + gate_load = () next_decoder_cache = None for decoder_layer in self.layers: @@ -1463,6 +1497,8 @@ class DeepseekModel(DeepseekPreTrainedModel): if output_attentions: all_self_attns += (layer_outputs[1],) + gate_load += (layer_outputs[-1],) + hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer @@ -1482,14 +1518,20 @@ class DeepseekModel(DeepseekPreTrainedModel): for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None ) - return BaseModelOutputWithPast( + return BaseMoEModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, + gate_load=gate_load, ) +@dataclass +class MoECausalLMOutputWithPast(CausalLMOutputWithPast): + gate_load: Optional[torch.Tensor] = None + + class DeepseekForCausalLM(DeepseekPreTrainedModel): _tied_weights_keys = ["lm_head.weight"] @@ -1620,12 +1662,13 @@ class DeepseekForCausalLM(DeepseekPreTrainedModel): output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( + return MoECausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, + gate_load=outputs.gate_load, ) def prepare_inputs_for_generation( diff --git a/src/utils/config.py b/src/utils/config.py index 3ea5283..d4060d9 100644 --- a/src/utils/config.py +++ b/src/utils/config.py @@ -6,6 +6,7 @@ import torch import transformers from src.utils.io import load_json +from src.data import get_uniform_sampling_ratio @dataclass @@ -33,7 +34,9 @@ class ModelArguments: ) attn_impl: str = field( default="flash_attention_2", - metadata={"help": "attention implementation, choice from [eager, flash_attention_2, sdpa] (default: `flash_attention_2`)"} + metadata={ + "help": "attention implementation, choice from [eager, flash_attention_2, sdpa] (default: `flash_attention_2`)" + }, ) def __post_init__(self): @@ -56,6 +59,18 @@ class DataArguments: default="data/merged", metadata={"help": "Path to dataset directory or a single jsonl file"}, ) + prob_map: str = field( + default=None, + metadata={"help": "Path to the probability map file"}, + ) + + def __post_init__(self): + if self.prob_map is not None: + if not pathlib.Path(self.prob_map).exists(): + raise ValueError(f"Probability map file {self.prob_map} not found") + self.prob_map = load_json(self.prob_map) + else: + self.prob_map = get_uniform_sampling_ratio(self.dataset_dir_or_path) @dataclass @@ -70,9 +85,7 @@ class TrainingArguments(transformers.TrainingArguments): ) max_eval_steps_per_type: int = field( default=10, - metadata={ - "help": "Maximum number of steps to perform during evaluation." - }, + metadata={"help": "Maximum number of steps to perform during evaluation."}, ) dynamic_sampling_sim_type: Literal["cos", "l2"] = field( default="l2", @@ -88,7 +101,5 @@ class TrainingArguments(transformers.TrainingArguments): ) freeze_gate: bool = field( default=False, - metadata={ - "help": "Whether to freeze the gate during training." - }, + metadata={"help": "Whether to freeze the gate during training."}, ) diff --git a/src/utils/visualization.py b/src/utils/visualization.py index 794f6c8..02bd236 100644 --- a/src/utils/visualization.py +++ b/src/utils/visualization.py @@ -180,6 +180,86 @@ def gate_load_stats(model_dir, data_dir, result_dir, update_strategy: str = "cos ) +def sampling_info_stats(filepath: str, data_type: str, output_dir: str): + from pathlib import Path + import numpy as np + from src.utils.io import load_jsonlines + + Path(output_dir).mkdir(exist_ok=True, parents=True) + + data = load_jsonlines(filepath) + step2data = {ins["step"]: ins for ins in data} + + data_types = sorted(data[0]["old_prob_map"].keys()) + data_type_idx = data_types.index(data_type) + + probs = [] + loads = [] + sims = [] + steps = sorted(step2data.keys()) + for step in steps: + ins = step2data[step] + probs.append(ins["old_prob_map"][data_type]) + loads.append(ins["name2load"][data_type]) + sims.append(ins["sim"][data_type_idx]) + + # probs + fig = plt.figure() + ax = fig.add_subplot(111) + ax.plot(steps, probs) + ax.set_title(f"Sampling Probability of {data_type}") + ax.set_xlabel("step") + fig.savefig(f"{output_dir}/prob-{data_type}.png") + + # loads + def cv_square(data): + return np.var(data, axis=1) / (np.mean(data, axis=1)**2 + 1e-10) + + fig = plt.figure() + ax = fig.add_subplot(111) + ax.plot(steps, cv_square(loads)) + ax.set_title(f"cv(load)^2 of {data_type}") + ax.set_xlabel("step") + fig.savefig(f"{output_dir}/load_cv-{data_type}.png") + + # sims + fig = plt.figure() + ax = fig.add_subplot(111) + ax.plot(steps, np.mean(sims, axis=1)) + ax.set_title(f"Mean Similarities with {data_type}") + ax.set_xlabel("step") + fig.savefig(f"{output_dir}/sim-{data_type}.png") + + +def test_sampling_convergence(): + from collections import defaultdict + from src.callbacks import AdaptiveSamplingCallback + + # freeze gate + name2load = {"code": [0.1359794776119403, 0.1333115671641791, 0.12858208955223882, 0.10330223880597016, 0.12544776119402984, 0.12625932835820897, 0.12761194029850748, 0.11950559701492537], "orca": [0.1509941502743006, 0.11721425756978752, 0.1232988815809414, 0.12714439426545024, 0.11256554420634679, 0.14008274482465977, 0.11819552632376563, 0.11050450095474797], "math": [0.15956486572028086, 0.10727138452881943, 0.11506675888262392, 0.10958069091633744, 0.11805010139847842, 0.11915200393871546, 0.13648938539627462, 0.13482480921846976], "sharegpt": [0.15337086599959998, 0.11428233411553493, 0.12873151621889287, 0.1177436980734424, 0.11538123789498336, 0.13793986642403783, 0.12419686111124664, 0.10835362016226212]} # fmt: skip + # # dynamic + # name2load = {"code": [0.14031716417910448, 0.1310634328358209, 0.12651119402985075, 0.10993470149253731, 0.12196828358208955, 0.12552238805970148, 0.12791977611940297, 0.11676305970149255], "orca": [0.15106234655836084, 0.11803640166095838, 0.12349968175067437, 0.12884551268450883, 0.11344072985178673, 0.1383778377231534, 0.11733170672566907, 0.1094057830448883], "math": [0.16001617686708006, 0.10756444371505268, 0.11391210568886491, 0.114803005615014, 0.11676650216277679, 0.1177863481308685, 0.13630182751708533, 0.13284959030325763], "sharegpt": [0.15440024978412215, 0.113654214863131, 0.12914741653941664, 0.12104040941178769, 0.11470799162832905, 0.13593110446537907, 0.12316259873058931, 0.10795601457724527]} # fmt: skip + names = sorted(name2load.keys()) + callback = AdaptiveSamplingCallback() + callback.prob_map = {"code": 0.25, "math": 0.25, "orca": 0.25, "sharegpt": 0.25} + name2probs = defaultdict(list) + for _ in range(100): + for name in names: + name2probs[name].append(callback.prob_map[name]) + new_name2prob, _ = callback._update_prob_map(name2load) + callback.prob_map = new_name2prob + print(f"final prob_map: {callback.prob_map}") + + fig = plt.figure() + ax = fig.add_subplot(111) + for name in names: + ax.plot(name2probs[name], label=name) + ax.legend() + ax.set_title("Sampling Probability") + ax.set_xlabel("step") + fig.savefig("results/sampling_convergence.png") + + if __name__ == "__main__": # gate_load_stats( # "/mnt/petrelfs/zhutong/llama-moe-models/LLaMA-MoE-v1-3_5B-2_8-new", @@ -195,12 +275,12 @@ if __name__ == "__main__": # "results/gate_load_vis_llama_moe_2_8_orca_4clusters", # ) - gate_load_stats( - "/mnt/petrelfs/zhutong/llama-moe-models/LLaMA-MoE-v1-3_5B-2_8-new", - "data/four_types_mix/dev", - "results/debug", - update_strategy="l2", - ) + # gate_load_stats( + # "/mnt/petrelfs/zhutong/llama-moe-models/LLaMA-MoE-v1-3_5B-2_8-new", + # "data/four_types_mix/dev", + # "results/debug", + # update_strategy="l2", + # ) # gate_load_stats( # "/mnt/petrelfs/zhutong/llama-moe-models/LLaMA-MoE-v1-3_5B-2_8-new", @@ -227,3 +307,29 @@ if __name__ == "__main__": # "results/gate_load_vis_llama_moe_2_8_four_types_mix_l2", # update_strategy="l2" # ) + + # sampling_info_stats( + # "/mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048_dynamic_remove_padding_tokens/llama_moe_four_mix_wo_pad_freeze_gate/moe_sft-2491632/sampling_info/data.jsonl", + # "code", + # "results/sampling_info/llama_moe_four_mix_wo_pad_freeze_gate/code", + # ) + + # sampling_info_stats( + # "/mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048_dynamic_remove_padding_tokens/llama_moe_four_mix_wo_pad/moe_sft-2491633/sampling_info/data.jsonl", + # "code", + # "results/sampling_info/llama_moe_four_mix_wo_pad/code", + # ) + + # sampling_info_stats( + # "/mnt/petrelfs/zhutong/adaptive-sft-for-moe/outputs/len2048_dynamic_remove_padding_tokens/llama_moe_four_mix_wo_pad_freeze_gate_wo_gate_noise/moe_sft-2493315/sampling_info/data.jsonl", + # "code", + # "results/sampling_info/llama_moe_four_mix_wo_pad_freeze_gate_wo_gate_noise/code", + # ) + + # sampling_info_stats( + # "outputs/len2048_dynamic_remove_padding_tokens/llama_moe_four_mix_wo_pad_wo_gate_noise/moe_sft-2492650/sampling_info/data.jsonl", + # "code", + # "results/sampling_info/llama_moe_four_mix_wo_pad_wo_gate_noise/code", + # ) + + test_sampling_convergence()