diff --git a/.gitignore b/.gitignore index a5669eb..8f2624b 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ plots/ *.safetensors *tfevents* watcher_state.yaml +ds_config.json \ No newline at end of file diff --git a/README.md b/README.md index 4e5ac4d..77cba6c 100644 --- a/README.md +++ b/README.md @@ -19,9 +19,11 @@ WANDB_MODE=disabled run python hyperlora/intx_sft.py configs/context_numbers_128 ### Generate fineweb qa ```bash # this might take several days... -python process_fineweb.py -python generate_fw_qa.py -python post_process_fw_qa.py +# python process_fineweb.py +# python generate_fw_qa.py +# python post_process_fw_qa.py + +vllm_model=mistralai/Mistral-Small-24B-Instruct-2501 run python generate_fw_qa_vllm.py 00_* 2 ``` diff --git a/configs/pretrain_all_xl.yaml b/configs/pretrain_all_xl.yaml new file mode 100644 index 0000000..d027400 --- /dev/null +++ b/configs/pretrain_all_xl.yaml @@ -0,0 +1,61 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: meta-llama/Llama-3.2-1B-Instruct +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00002 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 5 +weight_decay: 0.01 +# +warmup_steps: 100 + +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 +# LoRA +lora_r: 8 +lora_dropout: 0.05 +target_modules: + - down_proj + - up_proj + +# data +train_ds_names: +- fw_qa_xl +- ctx_qa +- pwc +- hotpot_qa +- squad +- drop +- narrativeqa +- quoref +- ropes +- synthetic_convqa + +val_ds_names: +- fw_qa_xl +- ctx_qa +- pwc +- hotpot_qa +- squad + +load_best_model_at_end: true +metric_for_best_model: eval_pwc_loss diff --git a/generate_fw_qa.py b/generate_fw_qa.py index 9a2d2f3..da674f6 100644 --- a/generate_fw_qa.py +++ b/generate_fw_qa.py @@ -93,32 +93,34 @@ def parse_args() -> Namespace: if __name__ == "__main__": # Load the datasets args = parse_args() - # ds = load_dataset( - # "parquet", - # data_files=f"./data/raw_datasets/fineweb_sharded/{args.shard_pattern}.parquet", - # split="train", - # streaming=True, - # ) - # os.makedirs("openai_batches", exist_ok=True) - # lines = [] - # c = 0 - # for i, sample in tqdm(enumerate(iter(ds))): - # if len(lines) >= 20_000: - # with open(f"openai_batches/fineweb_qa_pairs_{c}.jsonl", "w") as f: - # for line in lines: - # f.write(json.dumps(line) + "\n") - # lines = [] - # c += 1 - # jsonl = get_json_request( - # f"{c}_{i}", sample["text"], args.n_qa_pairs, args.gpt_model_name - # ) - # lines.append(jsonl) - # + ds = load_dataset( + "parquet", + data_files=f"./data/raw_datasets/fineweb_sharded/{args.shard_pattern}.parquet", + split="train", + streaming=True, + ) + os.makedirs("openai_batches", exist_ok=True) + lines = [] + c = 0 + for i, sample in tqdm(enumerate(iter(ds))): + if len(lines) >= 20_000: + with open(f"openai_batches/fineweb_qa_pairs_{c}.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n") + lines = [] + c += 1 + jsonl = get_json_request( + f"{c}_{i}", sample["text"], args.n_qa_pairs, args.gpt_model_name + ) + lines.append(jsonl) + client = OpenAI() res_files = glob("openai_batches/fineweb_qa_pairs_*_res.jsonl") prompt_files = glob("openai_batches/fineweb_qa_pairs_*[!res].jsonl") - unprocessed_files = set(prompt_files) - set([f.replace("_res","") for f in res_files]) - for file in unprocessed_files: + unprocessed_files = set(prompt_files) - set( + [f.replace("_res", "") for f in res_files] + ) + for file in unprocessed_files: print(f"Submitting batch {file}") batch_input_file = client.files.create(file=open(file, "rb"), purpose="batch") batch_input_file_id = batch_input_file.id diff --git a/generate_fw_qa_vllm.py b/generate_fw_qa_vllm.py new file mode 100644 index 0000000..919bfe2 --- /dev/null +++ b/generate_fw_qa_vllm.py @@ -0,0 +1,173 @@ +import os +import sys +import random +import re +from glob import glob + +import pandas as pd +from datasets import Dataset +from vllm import LLM, SamplingParams +from tqdm import tqdm +from datasets import load_dataset +from openai import OpenAI + +SYSTEM_TEMPLATE = ( + "You are a creative and helpful assistant.\n" + "You are given a context and you need to generate questions and corresponding answers from the given context.\n" + "The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n" + "Do not halucinate and make up information." +) + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " + "information provided in the context, not general questions that suits any context.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from information present in given context.\n" + "2. Make sure the questions are clear and unambiguous.\n" + "3. Phrases like 'based on the provided context', 'according to the context', etc, are not allowed to appear in " + "the questions.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the context.\n" + "2. Do not just copy words from the context. Answer the question in your own words.\n" + "3. The answers should be detailed and comprehensive.\n\n" + "Response with {n_qa_pairs} question-answer pairs. Use simple words and please be clear.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." + "\n\n" + "### Context ###\n" + "{context}" +) + + +def get_prompt(context, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs) + return prompt + + +def postprocess_qa_pairs(res_txt: str): + """ + Postprocesses the QA pairs from the response text. + + Args: + res_txt: The response text. + n_qa_pairs: The number of QA pairs. + + Returns: + A tuple of two lists, the first containing the questions and the second containing the answers. + """ + # capture everything after each "Question {number}:" until "Answer" + q_pattern = r"Question \d+:(.*?)(?=Answer|$)" # thanks chatgpt + questions = re.findall(q_pattern, res_txt, flags=re.S) + + a_pattern = r"Answer \d+:(.*?)(?=Question|$)" # thanks chatgpt + answers = re.findall(a_pattern, res_txt, flags=re.S) + + if len(questions) != len(answers): + print(f"Warning---number of questions and answers do not match") + print(f"Number of questions: {len(questions)}") + print(f"Number of answers: {len(answers)}") + + out_q = [] + out_a = [] + if (len(questions) > 0) and (len(answers) > 0): + for i in range(min(len(questions), len(answers))): + out_q.append(questions[i].strip()) + out_a.append(answers[i].strip()) + + return out_q, out_a + + +if __name__ == "__main__": + # api_key = os.environ.get("vllm_api_key") + # vllm_port = os.environ.get("vllm_port") + # vllm_ip = os.environ.get("vllm_ip") # "172.16.0.62" + vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it" + print(f"Using model: {vllm_model}") + # print(f"Using API key: {api_key}") + # print(f"Using VLLM IP: {vllm_ip}") + # print(f"Using VLLM port: {vllm_port}") + # client = OpenAI( + # base_url=f"http://{vllm_ip}:{vllm_port}/v1", + # api_key=api_key, + # ) + llm = LLM( + model=vllm_model, + tokenizer_mode="mistral", + config_format="mistral", + load_format="mistral", + enable_prefix_caching=True, + # enable_chunked_prefill=True, + ) + tokenizer = llm.get_tokenizer() + shard_pattern = sys.argv[1] + n_qa_pairs = int(sys.argv[2]) + + for path in glob(f"./data/raw_datasets/fineweb_sharded/{shard_pattern}.parquet"): + ds = load_dataset( + "parquet", + data_files=path, + split="train", + streaming=True, + ) + + ctxs = [sample["text"] for sample in iter(ds)] + messages = [ + [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": get_prompt(ctx, n_qa_pairs)}, + ] + for ctx in ctxs + ] + + print(f"Generating from {len(messages)} contexts") + completions = llm.chat( + messages, + sampling_params=SamplingParams( + max_tokens=2048, + temperature=1.0, + frequency_penalty=0.2, + ), + ) + samples = [] + for ctx, completion in zip(ctxs, completions): + questions, answers = postprocess_qa_pairs(completion.outputs[0].text) + for q, a in zip(questions, answers): + samples.append( + { + "context": ctx, + "prompt": q, + "response": a, + } + ) + + print(f"Generated {len(samples)} samples") + random.shuffle(samples) + df = pd.DataFrame(samples) + ds = Dataset.from_pandas(df) + val_ds = ds.take(10) + ds = ds.skip(10) + + shard_name = path.split("/")[-1].split(".")[0] + ds.to_parquet(f"data/raw_datasets/fw_qa_xl/{shard_name}.parquet") + val_ds.to_parquet(f"data/raw_datasets/fw_qa_xl/{shard_name}_val.parquet") + print(f"Saved to data/raw_datasets/fw_qa_xl/{shard_name}.parquet") + print(f"Saved to data/raw_datasets/fw_qa_xl/{shard_name}_val.parquet") + + # for i, sample in tqdm(enumerate(iter(ds))): + # completion = client.chat.completions.create( + # model=vllm_model, + # messages=[ + # {"role": "system", "content": SYSTEM_TEMPLATE}, + # {"role": "user", "content": get_prompt(sample["text"], n_qa_pairs)}, + # ], + # extra_body={"temperature": 1.0, "frequency_penalty": 0.2}, + # ) + + # print(completion.choices[0].message) + # if i >= 10: + # break diff --git a/generate_qa_parallel_vllm.py b/generate_qa_parallel_vllm.py new file mode 100644 index 0000000..3127687 --- /dev/null +++ b/generate_qa_parallel_vllm.py @@ -0,0 +1,516 @@ +# from https://github.com/openai/openai-cookbook/blob/main/examples/api_request_parallel_processor.py +""" +API REQUEST PARALLEL PROCESSOR + +Using the OpenAI API to process lots of text quickly takes some care. +If you trickle in a million API requests one by one, they'll take days to complete. +If you flood a million API requests in parallel, they'll exceed the rate limits and fail with errors. +To maximize throughput, parallel requests need to be throttled to stay under rate limits. + +This script parallelizes requests to the OpenAI API while throttling to stay under rate limits. + +Features: +- Streams requests from file, to avoid running out of memory for giant jobs +- Makes requests concurrently, to maximize throughput +- Throttles request and token usage, to stay under rate limits +- Retries failed requests up to {max_attempts} times, to avoid missing data +- Logs errors, to diagnose problems with requests + +Example command to call script: +``` +python examples/api_request_parallel_processor.py \ + --requests_filepath examples/data/example_requests_to_parallel_process.jsonl \ + --save_filepath examples/data/example_requests_to_parallel_process_results.jsonl \ + --request_url https://api.openai.com/v1/embeddings \ + --max_requests_per_minute 1500 \ + --max_tokens_per_minute 6250000 \ + --token_encoding_name cl100k_base \ + --max_attempts 5 \ + --logging_level 20 +``` + +Inputs: +- requests_filepath : str + - path to the file containing the requests to be processed + - file should be a jsonl file, where each line is a json object with API parameters and an optional metadata field + - e.g., {"model": "text-embedding-3-small", "input": "embed me", "metadata": {"row_id": 1}} + - as with all jsonl files, take care that newlines in the content are properly escaped (json.dumps does this automatically) + - an example file is provided at examples/data/example_requests_to_parallel_process.jsonl + - the code to generate the example file is appended to the bottom of this script +- save_filepath : str, optional + - path to the file where the results will be saved + - file will be a jsonl file, where each line is an array with the original request plus the API response + - e.g., [{"model": "text-embedding-3-small", "input": "embed me"}, {...}] + - if omitted, results will be saved to {requests_filename}_results.jsonl +- request_url : str, optional + - URL of the API endpoint to call + - if omitted, will default to "https://api.openai.com/v1/chat/completions" +- api_key : str, optional + - API key to use + - if omitted, the script will attempt to read it from an environment variable {os.getenv("OPENAI_API_KEY")} +- max_requests_per_minute : float, optional + - target number of requests to make per minute (will make less if limited by tokens) + - leave headroom by setting this to 50% or 75% of your limit + - if requests are limiting you, try batching multiple embeddings or completions into one request + - if omitted, will default to 1,500 +- max_tokens_per_minute : float, optional + - target number of tokens to use per minute (will use less if limited by requests) + - leave headroom by setting this to 50% or 75% of your limit + - if omitted, will default to 125,000 +- token_encoding_name : str, optional + - name of the token encoding used, as defined in the `tiktoken` package + - if omitted, will default to "o200k_base" (used by `gpt-4o-mini`) + - see https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken +- max_attempts : int, optional + - number of times to retry a failed request before giving up + - if omitted, will default to 5 +- logging_level : int, optional + - level of logging to use; higher numbers will log fewer messages + - 40 = ERROR; will log only when requests fail after all retries + - 30 = WARNING; will log when requests his rate limits or other errors + - 20 = INFO; will log when requests start and the status at finish + - 10 = DEBUG; will log various things as the loop runs to see when they occur + - if omitted, will default to 20 (INFO). + +The script is structured as follows: + - Imports + - Define main() + - Initialize things + - In main loop: + - Get next request if one is not already waiting for capacity + - Update available token & request capacity + - If enough capacity available, call API + - The loop pauses if a rate limit error is hit + - The loop breaks when no tasks remain + - Define dataclasses + - StatusTracker (stores script metadata counters; only one instance is created) + - APIRequest (stores API inputs, outputs, metadata; one method to call API) + - Define functions + - api_endpoint_from_url (extracts API endpoint from request URL) + - append_to_jsonl (writes to results file) + - num_tokens_consumed_from_request (bigger function to infer token usage from request) + - task_id_generator_function (yields 0, 1, 2, ...) + - Run main() +""" + +# imports +import aiohttp # for making API calls concurrently +import argparse # for running script from command line +import asyncio # for running API calls concurrently +import json # for saving results to a jsonl file +import logging # for logging rate limit warnings and other messages +import os # for reading API key +import re # for matching endpoint from request URL +import tiktoken # for counting tokens +import time # for sleeping after rate limit is hit +from dataclasses import ( + dataclass, + field, +) # for storing API inputs, outputs, and metadata + + +async def process_api_requests_from_file( + requests_filepath: str, + save_filepath: str, + request_url: str, + api_key: str, + max_requests_per_minute: float, + max_tokens_per_minute: float, + token_encoding_name: str, + max_attempts: int, + logging_level: int, +): + """Processes API requests in parallel, throttling to stay under rate limits.""" + # constants + seconds_to_pause_after_rate_limit_error = 15 + seconds_to_sleep_each_loop = ( + 0.01 # 10 ms limits max throughput to 100 requests per second + ) + + # initialize logging + logging.basicConfig(level=logging_level) + logging.debug(f"Logging initialized at level {logging_level}") + + # infer API endpoint and construct request header + api_endpoint = api_endpoint_from_url(request_url) + request_header = {"Authorization": f"Bearer {api_key}"} + # use api-key header for Azure deployments + if "/deployments" in request_url: + request_header = {"api-key": f"{api_key}"} + + # initialize trackers + queue_of_requests_to_retry = asyncio.Queue() + task_id_generator = ( + task_id_generator_function() + ) # generates integer IDs of 0, 1, 2, ... + status_tracker = ( + StatusTracker() + ) # single instance to track a collection of variables + next_request = None # variable to hold the next request to call + + # initialize available capacity counts + available_request_capacity = max_requests_per_minute + available_token_capacity = max_tokens_per_minute + last_update_time = time.time() + + # initialize flags + file_not_finished = True # after file is empty, we'll skip reading it + logging.debug(f"Initialization complete.") + + # initialize file reading + with open(requests_filepath) as file: + # `requests` will provide requests one at a time + requests = file.__iter__() + logging.debug(f"File opened. Entering main loop") + async with aiohttp.ClientSession() as session: # Initialize ClientSession here + while True: + # get next request (if one is not already waiting for capacity) + if next_request is None: + if not queue_of_requests_to_retry.empty(): + next_request = queue_of_requests_to_retry.get_nowait() + logging.debug( + f"Retrying request {next_request.task_id}: {next_request}" + ) + elif file_not_finished: + try: + # get new request + request_json = json.loads(next(requests)) + next_request = APIRequest( + task_id=next(task_id_generator), + request_json=request_json, + token_consumption=num_tokens_consumed_from_request( + request_json, api_endpoint, token_encoding_name + ), + attempts_left=max_attempts, + metadata=request_json.pop("metadata", None), + ) + status_tracker.num_tasks_started += 1 + status_tracker.num_tasks_in_progress += 1 + logging.debug( + f"Reading request {next_request.task_id}: {next_request}" + ) + except StopIteration: + # if file runs out, set flag to stop reading it + logging.debug("Read file exhausted") + file_not_finished = False + + # update available capacity + current_time = time.time() + seconds_since_update = current_time - last_update_time + available_request_capacity = min( + available_request_capacity + + max_requests_per_minute * seconds_since_update / 60.0, + max_requests_per_minute, + ) + available_token_capacity = min( + available_token_capacity + + max_tokens_per_minute * seconds_since_update / 60.0, + max_tokens_per_minute, + ) + last_update_time = current_time + + # if enough capacity available, call API + if next_request: + next_request_tokens = next_request.token_consumption + if ( + available_request_capacity >= 1 + and available_token_capacity >= next_request_tokens + ): + # update counters + available_request_capacity -= 1 + available_token_capacity -= next_request_tokens + next_request.attempts_left -= 1 + + # call API + asyncio.create_task( + next_request.call_api( + session=session, + request_url=request_url, + request_header=request_header, + retry_queue=queue_of_requests_to_retry, + save_filepath=save_filepath, + status_tracker=status_tracker, + ) + ) + next_request = None # reset next_request to empty + + # if all tasks are finished, break + if status_tracker.num_tasks_in_progress == 0: + break + + # main loop sleeps briefly so concurrent tasks can run + await asyncio.sleep(seconds_to_sleep_each_loop) + + # if a rate limit error was hit recently, pause to cool down + seconds_since_rate_limit_error = ( + time.time() - status_tracker.time_of_last_rate_limit_error + ) + if ( + seconds_since_rate_limit_error + < seconds_to_pause_after_rate_limit_error + ): + remaining_seconds_to_pause = ( + seconds_to_pause_after_rate_limit_error + - seconds_since_rate_limit_error + ) + await asyncio.sleep(remaining_seconds_to_pause) + # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago + logging.warn( + f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}" + ) + + # after finishing, log final status + logging.info( + f"""Parallel processing complete. Results saved to {save_filepath}""" + ) + if status_tracker.num_tasks_failed > 0: + logging.warning( + f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed. Errors logged to {save_filepath}." + ) + if status_tracker.num_rate_limit_errors > 0: + logging.warning( + f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate." + ) + + +# dataclasses + + +@dataclass +class StatusTracker: + """Stores metadata about the script's progress. Only one instance is created.""" + + num_tasks_started: int = 0 + num_tasks_in_progress: int = 0 # script ends when this reaches 0 + num_tasks_succeeded: int = 0 + num_tasks_failed: int = 0 + num_rate_limit_errors: int = 0 + num_api_errors: int = 0 # excluding rate limit errors, counted above + num_other_errors: int = 0 + time_of_last_rate_limit_error: int = 0 # used to cool off after hitting rate limits + + +@dataclass +class APIRequest: + """Stores an API request's inputs, outputs, and other metadata. Contains a method to make an API call.""" + + task_id: int + request_json: dict + token_consumption: int + attempts_left: int + metadata: dict + result: list = field(default_factory=list) + + async def call_api( + self, + session: aiohttp.ClientSession, + request_url: str, + request_header: dict, + retry_queue: asyncio.Queue, + save_filepath: str, + status_tracker: StatusTracker, + ): + """Calls the OpenAI API and saves results.""" + logging.info(f"Starting request #{self.task_id}") + error = None + try: + async with session.post( + url=request_url, headers=request_header, json=self.request_json + ) as response: + response = await response.json() + if "error" in response: + logging.warning( + f"Request {self.task_id} failed with error {response['error']}" + ) + status_tracker.num_api_errors += 1 + error = response + if "rate limit" in response["error"].get("message", "").lower(): + status_tracker.time_of_last_rate_limit_error = time.time() + status_tracker.num_rate_limit_errors += 1 + status_tracker.num_api_errors -= ( + 1 # rate limit errors are counted separately + ) + + except ( + Exception + ) as e: # catching naked exceptions is bad practice, but in this case we'll log & save them + logging.warning(f"Request {self.task_id} failed with Exception {e}") + status_tracker.num_other_errors += 1 + error = e + if error: + self.result.append(error) + if self.attempts_left: + retry_queue.put_nowait(self) + else: + logging.error( + f"Request {self.request_json} failed after all attempts. Saving errors: {self.result}" + ) + data = ( + [self.request_json, [str(e) for e in self.result], self.metadata] + if self.metadata + else [self.request_json, [str(e) for e in self.result]] + ) + append_to_jsonl(data, save_filepath) + status_tracker.num_tasks_in_progress -= 1 + status_tracker.num_tasks_failed += 1 + else: + data = ( + [self.request_json, response, self.metadata] + if self.metadata + else [self.request_json, response] + ) + append_to_jsonl(data, save_filepath) + status_tracker.num_tasks_in_progress -= 1 + status_tracker.num_tasks_succeeded += 1 + logging.debug(f"Request {self.task_id} saved to {save_filepath}") + + +# functions + + +def api_endpoint_from_url(request_url): + """Extract the API endpoint from the request URL.""" + match = re.search("^https://[^/]+/v\\d+/(.+)$", request_url) + if match is None: + # for Azure OpenAI deployment urls + match = re.search( + r"^https://[^/]+/openai/deployments/[^/]+/(.+?)(\?|$)", request_url + ) + return match[1] + + +def append_to_jsonl(data, filename: str) -> None: + """Append a json payload to the end of a jsonl file.""" + json_string = json.dumps(data) + with open(filename, "a") as f: + f.write(json_string + "\n") + + +def num_tokens_consumed_from_request( + request_json: dict, + api_endpoint: str, + token_encoding_name: str, +): + """Count the number of tokens in the request. Only supports completion and embedding requests.""" + encoding = tiktoken.get_encoding(token_encoding_name) + # if completions request, tokens = prompt + n * max_tokens + if api_endpoint.endswith("completions"): + max_tokens = request_json.get("max_tokens", 15) + n = request_json.get("n", 1) + completion_tokens = n * max_tokens + + # chat completions + if api_endpoint.startswith("chat/"): + num_tokens = 0 + for message in request_json["messages"]: + num_tokens += ( + 4 # every message follows {role/name}\n{content}\n + ) + for key, value in message.items(): + num_tokens += len(encoding.encode(value)) + if key == "name": # if there's a name, the role is omitted + num_tokens -= 1 # role is always required and always 1 token + num_tokens += 2 # every reply is primed with assistant + return num_tokens + completion_tokens + # normal completions + else: + prompt = request_json["prompt"] + if isinstance(prompt, str): # single prompt + prompt_tokens = len(encoding.encode(prompt)) + num_tokens = prompt_tokens + completion_tokens + return num_tokens + elif isinstance(prompt, list): # multiple prompts + prompt_tokens = sum([len(encoding.encode(p)) for p in prompt]) + num_tokens = prompt_tokens + completion_tokens * len(prompt) + return num_tokens + else: + raise TypeError( + 'Expecting either string or list of strings for "prompt" field in completion request' + ) + # if embeddings request, tokens = input tokens + elif api_endpoint == "embeddings": + input = request_json["input"] + if isinstance(input, str): # single input + num_tokens = len(encoding.encode(input)) + return num_tokens + elif isinstance(input, list): # multiple inputs + num_tokens = sum([len(encoding.encode(i)) for i in input]) + return num_tokens + else: + raise TypeError( + 'Expecting either string or list of strings for "inputs" field in embedding request' + ) + # more logic needed to support other API calls (e.g., edits, inserts, DALL-E) + else: + raise NotImplementedError( + f'API endpoint "{api_endpoint}" not implemented in this script' + ) + + +def task_id_generator_function(): + """Generate integers 0, 1, 2, and so on.""" + task_id = 0 + while True: + yield task_id + task_id += 1 + + +# run script + + +if __name__ == "__main__": + # parse command line arguments + parser = argparse.ArgumentParser() + parser.add_argument("--requests_filepath") + parser.add_argument("--save_filepath", default=None) + parser.add_argument( + "--request_url", default="https://api.openai.com/v1/chat/completions" + ) + parser.add_argument("--api_key", default=os.getenv("OPENAI_API_KEY")) + # tier 4 limit + parser.add_argument("--max_requests_per_minute", type=int, default=10_000 * 0.5) + parser.add_argument("--max_tokens_per_minute", type=int, default=10_000_000 * 0.5) + # + parser.add_argument("--token_encoding_name", default="o200k_base") + parser.add_argument("--max_attempts", type=int, default=5) + parser.add_argument("--logging_level", default=logging.INFO) + args = parser.parse_args() + + if args.save_filepath is None: + args.save_filepath = args.requests_filepath.replace(".jsonl", "_res.jsonl") + + # run script + asyncio.run( + process_api_requests_from_file( + requests_filepath=args.requests_filepath, + save_filepath=args.save_filepath, + request_url=args.request_url, + api_key=args.api_key, + max_requests_per_minute=float(args.max_requests_per_minute), + max_tokens_per_minute=float(args.max_tokens_per_minute), + token_encoding_name=args.token_encoding_name, + max_attempts=int(args.max_attempts), + logging_level=int(args.logging_level), + ) + ) + +""" +from openai import OpenAI +import os +api_key = os.environ.get("vllm_api_key") +vllm_port = os.environ.get("vllm_port") +vllm_ip = "172.16.0.62" +client = OpenAI( + base_url=f"http://{vllm_ip}:{vllm_port}/v1", + api_key=api_key, +) + +completion = client.chat.completions.create( + model="google/gemma-2-27b-it", + messages=[ + {"role": "user", "content": "Hello!"} + ] +) + +print(completion.choices[0].message) +""" diff --git a/hotpot_qa_lengths.png b/hotpot_qa_lengths.png deleted file mode 100644 index ea4cbd9..0000000 Binary files a/hotpot_qa_lengths.png and /dev/null differ diff --git a/hotpot_qa_lengths_filtered.png b/hotpot_qa_lengths_filtered.png deleted file mode 100644 index 63ddbf7..0000000 Binary files a/hotpot_qa_lengths_filtered.png and /dev/null differ diff --git a/intx_sft.py b/intx_sft.py index 9fed971..972e6d2 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -529,10 +529,9 @@ def main(): # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer # TODO: use packing with SFTTrainer - # HACK [local patch]: see transformers/trainer_seq2seq.py for supressing - # "Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead." # HACK [local patch]: deepspeed model loading problem (for resume training) # see https://github.com/microsoft/DeepSpeed/pull/6626/files + # /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora/lib/python3.10/site-packages/deepspeed/runtime/engine.py if training_args.use_liger_kernel and is_liger_kernel_available(): from liger_kernel.transformers import _apply_liger_kernel_to_instance @@ -593,4 +592,6 @@ if __name__ == "__main__": os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" if os.getenv("DEBUG", False): disable_caching() + # randomly sleep to avoid run_name collision + time.sleep(random.random() * 13) main() diff --git a/openphi_prog_lengths.png b/openphi_prog_lengths.png deleted file mode 100644 index 9c2f399..0000000 Binary files a/openphi_prog_lengths.png and /dev/null differ diff --git a/openphi_prog_lengths_filtered.png b/openphi_prog_lengths_filtered.png deleted file mode 100644 index c0e832a..0000000 Binary files a/openphi_prog_lengths_filtered.png and /dev/null differ diff --git a/process_fineweb.py b/process_fineweb.py index 5dda98d..58f03b8 100644 --- a/process_fineweb.py +++ b/process_fineweb.py @@ -22,7 +22,7 @@ if __name__ == "__main__": # https://github.com/huggingface/datasets/issues/7047#issuecomment-2233163406 num_shards_per_file = 16 sharded_fw_dir = "./data/raw_datasets/fineweb_sharded/" - output_path_template = f"{sharded_fw_dir}" + "/{index:05d}.parquet" + output_path_template = f"{sharded_fw_dir}" + "/{i:02d}_{idx:05d}.parquet" for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/10BT/*.parquet"))): # ~1M rows ~= 2GB mem required per file @@ -31,7 +31,9 @@ if __name__ == "__main__": print(f"Filtered ds size: {len(ds)}") ds = ds.shuffle(seed=42 + i) # take one shard (from 16 shards) per file - idx = random.sample(range(num_shards_per_file), 1)[0] - - shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=False) - shard.to_parquet(output_path_template.format(index=i)) + # idx = random.sample(range(num_shards_per_file), 1)[0] + # shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=False) + # shard.to_parquet(output_path_template.format(index=i)) + for idx in range(num_shards_per_file): + shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=True) + shard.to_parquet(output_path_template.format(i=i, idx=idx)) diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 91d182f..1aea333 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -133,8 +133,11 @@ class TrainingArguments(TrainingArguments): default=True, metadata={"help": "Whether to pin memory in data loaders or not."}, ) + # mem leak if use persistent workers + # https://github.com/pytorch/pytorch/issues/62066 + # https://github.com/huggingface/transformers/issues/30943 dataloader_persistent_workers: bool = field( - default=True, + default=False, metadata={ "help": "Whether to keep the workers alive after a dataset has been consumed once." }, @@ -156,11 +159,11 @@ class TrainingArguments(TrainingArguments): metadata={"help": "Adam beta 1."}, ) adam_beta2: float = field( - default=0.95, + default=0.999, metadata={"help": "Adam beta 2."}, ) adam_epsilon: float = field( - default=1e-6, + default=1e-8, metadata={"help": "Adam epsilon."}, ) lr_scheduler_type: str = field( @@ -200,7 +203,7 @@ class TrainingArguments(TrainingArguments): # metadata={"help": "Whether to load the best model at the end of training."}, # ) save_total_limit: int = field( - default=5, + default=2, metadata={"help": "Total number of checkpoints to save."}, ) save_strategy: str = field( diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 7edabe2..c7387f2 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -73,6 +73,18 @@ DS_KWARGS = { split="train", ), ), + "fw_qa_xl": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"), + split="train", + ), + ), "ctx_qa": dict( train=dict( path="parquet", diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index 24a4ab9..e75d47a 100644 --- a/src/ctx_to_lora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -341,8 +341,25 @@ def evaluate(checkpoint_path, args, split, generative): train=False, use_flash_attn=True, ) + if generative: + model = model.to(torch.bfloat16) # NOTE: there is still some randomness in the eval result # despite all the deterministic settings + if args.use_liger_kernel and is_liger_kernel_available(): + from liger_kernel.transformers import _apply_liger_kernel_to_instance + + if isinstance(model, ModulatedPretrainedModel): + print("Applying liger-kernel to ModulatedPretrainedModel") + if isinstance(model.base_model, PeftModel): + _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) + else: + _apply_liger_kernel_to_instance(model=model.base_model.model) + if ctx_name is not None: + print("Applying liger-kernel to ctx_encoder_model") + _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model) + elif isinstance(model, PeftModel): + print("Applying liger-kernel to PeftModel") + _apply_liger_kernel_to_instance(model=model.base_model.model) tokenizer = get_tokenizer(args.model_name_or_path, train=False) if tokenizer.pad_token_id is None: @@ -392,7 +409,7 @@ def evaluate(checkpoint_path, args, split, generative): eval_trainer_args["eval_strategy"] = "no" eval_trainer_args["save_strategy"] = "no" eval_trainer_args["overwrite_output_dir"] = True - eval_trainer_args["per_device_eval_batch_size"] = 4 + eval_trainer_args["per_device_eval_batch_size"] = 32 eval_trainer_args = Seq2SeqTrainingArguments( **eval_trainer_args, diff --git a/src/ctx_to_lora/modeling_idefics2.py b/src/ctx_to_lora/modeling_idefics2.py index c3d30b8..4f60c38 100644 --- a/src/ctx_to_lora/modeling_idefics2.py +++ b/src/ctx_to_lora/modeling_idefics2.py @@ -758,6 +758,7 @@ class Idefics2Perceiver(Idefics2PreTrainedModel): # outputs = self.decoder(latents) # else: # packed sequence + # always use position_ids for decoder latent_position_ids = torch.arange( self.encoder.n_latents, device=context.device ).unsqueeze(0) diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 27b7db6..04d7d5a 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -38,7 +38,7 @@ from transformers.models.perceiver.modeling_perceiver import ( PerceiverBasicDecoder, ) from transformers.modeling_outputs import ModelOutput - +from transformers.models.modernbert.modeling_modernbert import ModernBertModel from ctx_to_lora.configs import ( AggregatorArguments, HypernetArguments, @@ -372,6 +372,7 @@ def maybe_add_batch_dim(kwargs): if ( "attention_mask" in kwargs and kwargs["attention_mask"] is not None + and isinstance(kwargs["attention_mask"], torch.Tensor) and len(kwargs["attention_mask"].shape) == 1 ): kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0) @@ -405,14 +406,14 @@ class EarlyExit(nn.Module): # if len(kwargs["attention_mask"].shape) == 1: # kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0) - with ( - # early_exit(self.base_model, self.exit_layer), - maybe_add_batch_dim(kwargs) as (batched_input, batched_attn_mask), - ): - model_outputs = self.base_model(**kwargs) + # with ( + # # early_exit(self.base_model, self.exit_layer), + # maybe_add_batch_dim(kwargs) as (batched_input, batched_attn_mask), + # ): + model_outputs = self.base_model(**kwargs) - if batched_input: - model_outputs.last_hidden_state = model_outputs.last_hidden_state.squeeze(0) + # if batched_input: + # model_outputs.last_hidden_state = model_outputs.last_hidden_state.squeeze(0) return model_outputs.last_hidden_state @@ -767,13 +768,13 @@ class ModulatedPretrainedModel(nn.Module): get_base_model(encoder_model), self.ctx_encoder_args.layer_idx ) - def to(self, *args, **kwargs): - # workaround to avoid the hypernet being wrapped by DeepSpeed - self.base_model = self.base_model.to(*args, **kwargs) - self.ctx_encoder = self.ctx_encoder.to(*args, **kwargs) - # self.hypernet = self.hypernet.to(*args, **kwargs) - # self.hypernet.to(torch.float32) - return self + # def to(self, *args, **kwargs): + # # workaround to avoid the hypernet being wrapped by DeepSpeed + # self.base_model = self.base_model.to(*args, **kwargs) + # self.ctx_encoder = self.ctx_encoder.to(*args, **kwargs) + # # self.hypernet = self.hypernet.to(*args, **kwargs) + # # self.hypernet.to(torch.float32) + # return self # Delegate to base_model @property @@ -907,19 +908,46 @@ class ModulatedPretrainedModel(nn.Module): ctx_ids: Integer[Tensor, "bs ctx_len"], ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, ctx_position_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, - *args: Any, **kwargs: Any, ): with torch.no_grad(): - # TODO: for modernbert ctx_encoder pass - # `cu_seq_len` and `max_seq_len` to the forward call - ctx_features = self.ctx_encoder( + ctx_encoder_kwargs = dict( input_ids=ctx_ids, attention_mask=ctx_attn_mask, position_ids=ctx_position_ids, - *args, - **kwargs, ) + # TODO: for modernbert ctx_encoder pass + # `cu_seq_len` and `max_seq_len` to the forward call + if isinstance(self.ctx_encoder.base_model, ModernBertModel): + position_ids = ctx_position_ids.flatten() + indices = torch.arange( + position_ids.size(0), device=position_ids.device, dtype=torch.int32 + ) + # [bsz + 1] + cu_seqlens = torch.cat( + ( + indices[position_ids == 0], + torch.tensor( + position_ids.size(), + device=position_ids.device, + dtype=torch.int32, + ), + ) + ) + ctx_encoder_kwargs = dict( + input_ids=ctx_ids.squeeze(0), + cu_seqlens=cu_seqlens, + max_seqlen=position_ids.max() + 1, + attention_mask=-1, + seq_len=-1, + batch_size=-1, + ) + + ctx_features = self.ctx_encoder(**ctx_encoder_kwargs, **kwargs) + + if isinstance(self.ctx_encoder.base_model, ModernBertModel): + ctx_features = ctx_features.unsqueeze(0) + # print(f"padded ctx_features: {ctx_features.shape}") # if ctx_attn_mask is not None: # print(f"padded ctx_attn_mask: {ctx_attn_mask.shape}")