diff --git a/intx_sft.py b/intx_sft.py index 8781977..ede385f 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -55,6 +55,7 @@ from transformers import ( HfArgumentParser, set_seed, ) +from accelerate import Accelerator, DeepSpeedPlugin from peft import PeftModel from transformers.utils import is_liger_kernel_available from transformers.data import DataCollatorWithFlattening @@ -252,6 +253,8 @@ def main(): ] ) + accelerator = Accelerator(deepspeed_plugins=DeepSpeedPlugin("ds_config.json")) + device = accelerator.device set_seed(training_args.seed) checkpoint_dir = training_args.resume_from_checkpoint if checkpoint_dir and not os.path.isdir(checkpoint_dir): @@ -309,6 +312,7 @@ def main(): train=True, requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, peft_config=get_lora_config(model_name, **vars(lora_args)), + device=device, ) ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path if ctx_name is not None: @@ -376,6 +380,8 @@ def main(): logger.debug(model) log_num_train_params(model) + # model = accelerator.prepare(model) + ############ Dataset setup logger.info("Loading dataset...") @@ -405,12 +411,10 @@ def main(): if not ds_names: continue streaming = (split == "train") and data_args.streaming - tokenized_ds[split] = { - os.path.basename(ds_name): _get_tokenized_dataset( - ds_name, split, streaming=streaming - ) - for ds_name in ds_names - } + for ds_name in ds_names: + with accelerator.main_process_first(): + ds = _get_tokenized_dataset(ds_name, split, streaming=streaming) + tokenized_ds[split] = {os.path.basename(ds_name): ds} train_ds = tokenized_ds["train"] logging.info(f"train_ds: {train_ds}") @@ -575,7 +579,7 @@ if __name__ == "__main__": os.environ["WANDB_CONSOLE"] = "off" os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # os.environ["KMP_AFFINITY"] = "disabled" # fixing iterable dataset stuck - # os.environ["OMP_NUM_THREADS"] = "16" + os.environ["OMP_NUM_THREADS"] = "23" # os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = "137438953472" # 128 TB if os.getenv("DEBUG", False): disable_caching() diff --git a/scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus.sh b/scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus.sh new file mode 100644 index 0000000..a26d24a --- /dev/null +++ b/scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus.sh @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora_medium +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29564 intx_sft.py configs/pretrain_all_xl_3_medium.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=2 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True \ +--per_layer_processing=True diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index ffee7b1..99834d3 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -410,6 +410,7 @@ def _load_and_process_dataset( add_repeat_prompt: bool, streaming: bool, ds_path: str, + num_proc: int, ): try: ds_kwargs = get_ds_kwargs(ds_name, split) @@ -447,27 +448,27 @@ def _load_and_process_dataset( ds = ds.map( get_preprocessing_fn(ds_name), remove_columns=cols_to_remove, - num_proc=16, + num_proc=num_proc, ) # ds = ds.remove_columns(cols_to_remove) - ds = ds.filter(filter_none, batched=True, num_proc=16) - ds = ds.filter(filter_long_samples, batched=True, num_proc=16) + ds = ds.filter(filter_none, batched=True, num_proc=num_proc) + ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc) if split == "train": if add_negative_prompt: ds = ds.map( add_negative_prompt_fn, batched=True, batch_size=100_000, - num_proc=16, + num_proc=num_proc, ) if add_repeat_prompt and "context_numbers" not in ds_name: ds = ds.map( add_repeat_prompt_fn, batched=True, batch_size=100_000, - num_proc=16, + num_proc=num_proc, ) - ds.save_to_disk(ds_path, num_proc=16) + ds.save_to_disk(ds_path, num_proc=num_proc) return ds @@ -496,6 +497,7 @@ def get_tokenized_dataset( add_repeat_prompt=add_repeat_prompt, streaming=streaming, ) + num_proc = None if streaming and split == "train" else 8 ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest() ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" @@ -508,6 +510,7 @@ def get_tokenized_dataset( _load_and_process_dataset( **load_and_process_kwargs, ds_path=ds_path, + num_proc=num_proc, ) ds = datasets.load_from_disk(ds_path) @@ -521,6 +524,7 @@ def get_tokenized_dataset( need_ctx_ids, ds, set_format, + num_proc, ) return tokenized_ds @@ -535,6 +539,7 @@ def construct_and_tokenize_ctx_qa( need_ctx_ids, ds, set_format=None, + num_proc=None, ): kwargs = dict( tokenizer=repr(tokenizer), @@ -561,21 +566,21 @@ def construct_and_tokenize_ctx_qa( ds = ds.map( convert_ctx_prompt_response_to_messages, fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat}, - num_proc=16, + num_proc=num_proc, ) # add "chat" field ds = ds.map( get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), batched=True, batch_size=100_000, - num_proc=16, + num_proc=num_proc, ) # tokenize the chat + mask the assistant inputs ds = ds.filter( filter_long_chat, batched=True, batch_size=100_000, - num_proc=16, + num_proc=num_proc, ) # add "input_ids", "attention_mask", "labels" @@ -586,7 +591,7 @@ def construct_and_tokenize_ctx_qa( "mask_assistant_inputs": True, "tokenizer_kwargs": tokenizer_kwargs, }, - num_proc=16, + num_proc=num_proc, ) # for use_kl_loss, we need "chat_ids" and "chat_attn_mask" @@ -618,7 +623,7 @@ def construct_and_tokenize_ctx_qa( fn_kwargs={"tokenizer": ctx_tokenizer}, batched=True, batch_size=100_000, - num_proc=16, + num_proc=num_proc, ) tokenized_ds = tokenized_ds.remove_columns( @@ -630,7 +635,7 @@ def construct_and_tokenize_ctx_qa( # # the columns are unknown when using streaming dataset # tokenized_ds = tokenized_ds._resolve_features() # validate_columns(tokenized_ds) - tokenized_ds.save_to_disk(ds_path, num_proc=16) + tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) del tokenized_ds return datasets.load_from_disk(ds_path)