fw_qa_xl + modernbert

.gitignore

@@ -21,3 +21,4 @@ plots/
 *.safetensors
 *tfevents*
 watcher_state.yaml
+ds_config.json

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
 ```
 
 

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

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

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

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 <im_start>{role/name}\n{content}<im_end>\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 <im_start>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)
+"""

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()

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))

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(

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",

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,

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)

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}")