Skip to content

Instantly share code, notes, and snippets.

@galopyz

galopyz/sequence_packing_ipynb.ipynb Secret

Created September 7, 2025 00:39
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save galopyz/e9de8881eb3c526354942d45b7ab74ff to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save galopyz/e9de8881eb3c526354942d45b7ab74ff to your computer and use it in GitHub Desktop.
Download ZIP
My Dialog
Display the source blob
Display the rendered blob
Raw
sequence_packing_ipynb.ipynb
{
"cells": [
{
"cell_type": "markdown",
"id": "540a13f9",
"metadata": {
"input_tokens": 6
},
"source": [
"# Sequence packer"
]
},
{
"cell_type": "markdown",
"id": "b61f1929",
"metadata": {
"input_tokens": 15
},
"source": [
"Learning about how modernbert sequence packer works."
]
},
{
"cell_type": "markdown",
"id": "1fa17593",
"metadata": {
"input_tokens": 7
},
"source": [
"## Time and Batch implementation"
]
},
{
"cell_type": "markdown",
"id": "b62d863a",
"metadata": {
"input_tokens": 88,
"output_tokens": 271,
"solveit_ai": true
},
"source": [
"I have modernbert sequence packing code. I want to understand how it works by trying it out small bits at a time. However, I don't have composer installed on this machine and I lack space. Can I just go about without installing it? Or it is necessary to try out the code?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Looking at your code, I can see that the `Time` class from Composer is causing the error, but the good news is that **you can definitely explore this sequence packing code without installing Composer!**\n",
"\n",
"The sequence packer has minimal dependencies on Composer. Here's what I can see:\n",
"\n",
"1. The `Time` class is only used in the warmup schedulers for parsing time strings\n",
"2. The main sequence packing logic is completely independent\n",
"3. Most of the functionality uses standard libraries (numpy, torch, etc.)\n",
"\n",
"Would you like me to show you how to create a simple mock `Time` class so you can run the core functionality? Or would you prefer to start by exploring specific parts of the sequence packer that don't need the schedulers at all?\n",
"\n",
"What aspect of the sequence packing are you most curious about - the basic packing algorithm, the masking logic, or how it handles batching?"
]
},
{
"cell_type": "markdown",
"id": "4e610c88",
"metadata": {
"input_tokens": 48,
"output_tokens": 888,
"solveit_ai": true
},
"source": [
"Are you familiar with `Time`? If you are not confident,, you can use the web to learn more about it. Can you implement it plz?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! I found the Composer `Time` class. Time represents static durations of training time in terms of a TimeUnit enum. [^1] [^2] Let me search for more details about the TimeUnit enum and the `from_timestring` method.Great! Now I understand the Composer `Time` class better. Time represents static durations of training time in terms of a TimeUnit enum [^3] [^4] and A time string is a numerical value followed by the value of a TimeUnit enum. For example: Time.from_timestring(\"5ep\") describes 5 epochs [^5].\n",
"\n",
"Based on what I found, here's a minimal implementation of the `Time` class that should work for your sequence packer:\n",
"\n",
"```python\n",
"from enum import Enum\n",
"\n",
"class TimeUnit(Enum):\n",
" EPOCH = \"ep\"\n",
" BATCH = \"ba\" \n",
" TOKEN = \"tok\"\n",
" SAMPLE = \"sp\"\n",
" DURATION = \"dur\"\n",
"\n",
"class Time:\n",
" def __init__(self, value, unit):\n",
" self.value = value\n",
" self.unit = unit if isinstance(unit, TimeUnit) else TimeUnit(unit)\n",
" \n",
" @classmethod\n",
" def from_timestring(cls, timestring):\n",
" # Simple parser for strings like \"1000tok\", \"5ep\", etc.\n",
" import re\n",
" match = re.match(r'(\\d+(?:\\.\\d+)?)([a-z]+)', timestring)\n",
" if not match:\n",
" raise ValueError(f\"Invalid time string: {timestring}\")\n",
" value, unit_str = match.groups()\n",
" return cls(float(value), unit_str)\n",
"```\n",
"\n",
"Would you like to try this implementation, or would you prefer to start exploring the sequence packer without the warmup schedulers first?\n",
"\n",
"[^1]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/generated/composer.Time.html\n",
"\t\"Time represents static durations of training time in terms of a TimeUnit enum. \"\n",
"\n",
"[^2]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/generated/composer.Time.html\n",
"\t\"Time represents static durations of training time in terms of a TimeUnit enum. \"\n",
"\n",
"[^3]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/generated/composer.Time.html\n",
"\t\"Time represents static durations of training time in terms of a TimeUnit enum. \"\n",
"\n",
"[^4]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/generated/composer.Time.html\n",
"\t\"Time represents static durations of training time in terms of a TimeUnit enum. \"\n",
"\n",
"[^5]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/generated/composer.Time.html\n",
"\t\"A time string is a numerical value followed by the value of a TimeUnit enum. For example: >>> Time.from_timestring(\\\"5ep\\\") # describes 5 epochs. \""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a1a7426c",
"metadata": {
"input_tokens": 270
},
"outputs": [],
"source": [
"from enum import Enum\n",
"\n",
"class TimeUnit(Enum):\n",
" EPOCH = \"ep\"\n",
" BATCH = \"ba\" \n",
" TOKEN = \"tok\"\n",
" SAMPLE = \"sp\"\n",
" DURATION = \"dur\"\n",
"\n",
"class Time:\n",
" def __init__(self, value, unit):\n",
" self.value = value\n",
" self.unit = unit if isinstance(unit, TimeUnit) else TimeUnit(unit)\n",
" \n",
" @classmethod\n",
" def from_timestring(cls, timestring):\n",
" # Simple parser for strings like \"1000tok\", \"5ep\", etc.\n",
" import re\n",
" match = re.match(r'(\\d+(?:\\.\\d+)?)([a-z]+)', timestring)\n",
" if not match:\n",
" raise ValueError(f\"Invalid time string: {timestring}\")\n",
" value, unit_str = match.groups()\n",
" return cls(float(value), unit_str)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "76b368f4",
"metadata": {
"input_tokens": 22,
"output_tokens": 52
},
"outputs": [
{
"data": {
"text/plain": [
"(5.0, <TimeUnit.EPOCH: 'ep'>)"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"t = Time.from_timestring('5ep')\r\n",
"t.value, t.unit"
]
},
{
"cell_type": "markdown",
"id": "a842de1b",
"metadata": {
"input_tokens": 22,
"output_tokens": 406,
"solveit_ai": true
},
"source": [
"Thanks. Can you od the same thing with `Batch` now? \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Looking at the search results, I can see that A batch of data can be represented in several formats, depending on the application. [^1] and Can be any abritrary type that user creates, but we assume some sort of sequence (list, tuple, tensor, array), mapping (dictionary), or attribute store (object with data members, namedtuple). [^2]\n",
"\n",
"Based on this, `Batch` in Composer is just a type alias for `Any`. Let me create a simple implementation:\n",
"\n",
"```python\n",
"from typing import Any\n",
"\n",
"Batch = Any\n",
"```\n",
"\n",
"This should work for your sequence packer code since it just represents a flexible batch type that can hold different data structures like dictionaries, tensors, etc.\n",
"\n",
"Would you like to try adding this to your code and see if it resolves the error?\n",
"\n",
"[^1]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/composer.core.types.html\n",
"\t\"A batch of data can be represented in several formats, depending on the application.\"\n",
"\n",
"[^2]: https://docs.mosaicml.com/projects/composer/en/stable/api_reference/generated/composer.utils.batch_set.html\n",
"\t\"Can be any abritrary type that user creates, but we assume some sort of sequence (list, tuple, tensor, array), mapping (dictionary), or attribute stor...\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "71ef9fa9",
"metadata": {
"input_tokens": 12
},
"outputs": [],
"source": [
"from typing import Any\n",
"\n",
"Batch = Any"
]
},
{
"cell_type": "markdown",
"id": "bc95fea0",
"metadata": {
"heading_collapsed": true,
"input_tokens": 3
},
"source": [
"## Code"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2f423015",
"metadata": {
"hidden": 1,
"hide_input": true,
"input_tokens": 10692
},
"outputs": [],
"source": [
"# Copyright 2024 onwards Answer.AI, LightOn, and contributors\r\n",
"# License: Apache-2.0\r\n",
"\r\n",
"import math\r\n",
"import threading\r\n",
"import time\r\n",
"from abc import ABC, abstractmethod\r\n",
"from collections import deque\r\n",
"from enum import Enum\r\n",
"from typing import Any, Generic, Iterable, NamedTuple, Optional, Sequence, TypeVar, Union\r\n",
"\r\n",
"import numpy as np\r\n",
"import torch\r\n",
"# from composer.core import Time\r\n",
"# from composer.core.types import Batch\r\n",
"from numba import njit\r\n",
"\r\n",
"\r\n",
"class MaskingSchedule(str, Enum):\r\n",
" LINEAR = \"linear\"\r\n",
" STAIR = \"stair\"\r\n",
"\r\n",
"\r\n",
"class BatchSizeWarmupScheduler:\r\n",
" def __init__(\r\n",
" self,\r\n",
" min_batch_size: int,\r\n",
" max_batch_size: int,\r\n",
" warmup_tokens: Union[str, Time, int],\r\n",
" world_size: int,\r\n",
" ):\r\n",
" self.min_batch_size = min_batch_size\r\n",
" self.max_batch_size = max_batch_size\r\n",
"\r\n",
" if isinstance(warmup_tokens, str):\r\n",
" self.warmup_tokens = Time.from_timestring(warmup_tokens).value\r\n",
" elif isinstance(warmup_tokens, Time):\r\n",
" self.warmup_tokens = warmup_tokens.value\r\n",
" else:\r\n",
" self.warmup_tokens = warmup_tokens\r\n",
" self.warmup_tokens = math.ceil(self.warmup_tokens / world_size)\r\n",
" self._step_thresholds = self._calculate_step_thresholds()\r\n",
"\r\n",
" def _calculate_step_thresholds(self):\r\n",
" total_batch_sizes = sum(range(self.min_batch_size, self.max_batch_size))\r\n",
" steps_per_unit = self.warmup_tokens / total_batch_sizes\r\n",
"\r\n",
" thresholds = []\r\n",
" cumsum = 0\r\n",
" for batch_size in range(self.min_batch_size, self.max_batch_size):\r\n",
" cumsum += batch_size\r\n",
" steps = math.ceil(steps_per_unit * cumsum)\r\n",
" thresholds.append(steps)\r\n",
" return thresholds\r\n",
"\r\n",
" def __call__(self, current_step: int) -> int:\r\n",
" if current_step >= self.warmup_tokens:\r\n",
" return self.max_batch_size\r\n",
"\r\n",
" for i, threshold in enumerate(self._step_thresholds):\r\n",
" if current_step < threshold:\r\n",
" return self.min_batch_size + i\r\n",
"\r\n",
" # should never hit this, but just in case\r\n",
" return self.max_batch_size\r\n",
"\r\n",
"\r\n",
"class MaskingRatioCooldownScheduler:\r\n",
" def __init__(\r\n",
" self,\r\n",
" initial_masking_ratio: float,\r\n",
" final_masking_ratio: float,\r\n",
" cooldown_tokens: Union[str, Time, int],\r\n",
" cooldown_start_tokens: Union[str, Time, int] = 0,\r\n",
" schedule_type: MaskingSchedule = MaskingSchedule.LINEAR,\r\n",
" masking_ratio_step_size: Optional[float] = None,\r\n",
" world_size: int = 1,\r\n",
" ):\r\n",
" self.initial_masking_ratio = initial_masking_ratio\r\n",
" self.final_masking_ratio = final_masking_ratio\r\n",
" self.schedule_type = schedule_type\r\n",
" self.masking_ratio_step_size = masking_ratio_step_size\r\n",
" self.num_cooldown_steps = None\r\n",
"\r\n",
" if self.schedule_type == MaskingSchedule.STAIR:\r\n",
" if self.masking_ratio_step_size is None:\r\n",
" raise ValueError(\"`masking_ratio_step_size` must be provided for 'stair schedule.\")\r\n",
" if self.masking_ratio_step_size <= 0:\r\n",
" raise ValueError(\"`masking_ratio_step_size` must be positive.\")\r\n",
" if self.initial_masking_ratio > self.final_masking_ratio:\r\n",
" self.num_cooldown_steps = math.ceil(\r\n",
" (self.initial_masking_ratio - self.final_masking_ratio) / self.masking_ratio_step_size\r\n",
" )\r\n",
" else:\r\n",
" self.num_cooldown_steps = 0\r\n",
"\r\n",
" if isinstance(cooldown_start_tokens, str):\r\n",
" self.cooldown_start_tokens = Time.from_timestring(cooldown_start_tokens).value\r\n",
" elif isinstance(cooldown_start_tokens, Time):\r\n",
" self.cooldown_start_tokens = cooldown_start_tokens.value\r\n",
" else:\r\n",
" self.cooldown_start_tokens = cooldown_start_tokens\r\n",
" self.cooldown_start_tokens = math.ceil(self.cooldown_start_tokens / world_size)\r\n",
"\r\n",
" if isinstance(cooldown_tokens, str):\r\n",
" self.cooldown_tokens = Time.from_timestring(cooldown_tokens).value\r\n",
" elif isinstance(cooldown_tokens, Time):\r\n",
" self.cooldown_tokens = cooldown_tokens.value\r\n",
" else:\r\n",
" self.cooldown_tokens = cooldown_tokens\r\n",
" self.cooldown_tokens = math.ceil(self.cooldown_tokens / world_size)\r\n",
"\r\n",
" if self.cooldown_tokens < 1:\r\n",
" raise ValueError(\"`cooldown_tokens` divided by world_size must be greater than 1.\")\r\n",
"\r\n",
" self.cooldown_end_tokens = self.cooldown_start_tokens + self.cooldown_tokens\r\n",
"\r\n",
" def __call__(self, current_tokens: int) -> float:\r\n",
" if current_tokens < self.cooldown_start_tokens:\r\n",
" return self.initial_masking_ratio\r\n",
" if current_tokens >= self.cooldown_end_tokens:\r\n",
" return self.final_masking_ratio\r\n",
"\r\n",
" progress = (current_tokens - self.cooldown_start_tokens) / self.cooldown_tokens\r\n",
"\r\n",
" if self.schedule_type == MaskingSchedule.LINEAR:\r\n",
" current_masking_ratio = self.initial_masking_ratio - progress * (\r\n",
" self.initial_masking_ratio - self.final_masking_ratio\r\n",
" )\r\n",
" elif self.schedule_type == MaskingSchedule.STAIR:\r\n",
" if self.num_cooldown_steps is None or self.masking_ratio_step_size is None:\r\n",
" raise ValueError(\"`num_cooldown_steps` and `masking_ratio_step_size` must be set for stair.\")\r\n",
" step = math.floor(progress * self.num_cooldown_steps)\r\n",
" current_masking_ratio = self.initial_masking_ratio - step * self.masking_ratio_step_size\r\n",
" else:\r\n",
" raise ValueError(f\"Unknown schedule type: {self.schedule_type}\")\r\n",
" return max(current_masking_ratio, self.final_masking_ratio)\r\n",
"\r\n",
"\r\n",
"class SequencePackerBatchOutputTuple(NamedTuple):\r\n",
" masked_pseqs: torch.Tensor\r\n",
" labels: Optional[torch.Tensor]\r\n",
" cu_seq_lens: list[torch.Tensor]\r\n",
" max_cu_seq_len: list[torch.Tensor]\r\n",
"\r\n",
"\r\n",
"class SequencePacker(ABC):\r\n",
" def __init__(\r\n",
" self,\r\n",
" # params defining the incoming batches of seqs\r\n",
" src_iterable: Iterable[list[list[int]]],\r\n",
" src_batch_size: int,\r\n",
" src_max_seq_len: int,\r\n",
" # params defining outgoing batches of pseqs\r\n",
" out_batch_size: int,\r\n",
" out_pseq_len: int,\r\n",
" # params defining internal behavior\r\n",
" buffer_size: int,\r\n",
" pad_token_id: int = -1,\r\n",
" mask_token_id: int = 0,\r\n",
" ignore_token_id: int = -100,\r\n",
" mask_prob: float = 0.3,\r\n",
" seed=42,\r\n",
" suppress_masking: bool = False,\r\n",
" batch_size_warmup_min_size: Optional[int] = None,\r\n",
" batch_size_warmup_tokens: Optional[Union[str, Time]] = None,\r\n",
" masking_ratio_cooldown_tokens: Optional[Union[str, Time]] = None,\r\n",
" masking_ratio_cooldown_start_tokens: Optional[Union[str, Time]] = None,\r\n",
" final_masking_ratio: Optional[float] = None,\r\n",
" masking_cooldown_schedule_type: Union[str, MaskingSchedule] = \"linear\",\r\n",
" masking_ratio_step_size: Optional[float] = None,\r\n",
" world_size: int = 1,\r\n",
" ):\r\n",
" \"\"\"\r\n",
" Takes batches of unpacked, unpadded sequences (seqs) to batches of packed and padded sequences (pseqs).\r\n",
"\r\n",
" Every input batch must be a list[list[int]], a list of variable-length sequences of tokens.\r\n",
"\r\n",
" Every output batch is a tuple (masked_inputs:Tensor, labels:Tensor, seq_starts_and_end:list).\r\n",
"\r\n",
" It performs this streamwise, taking an iterable as the source of incoming batches, and\r\n",
" presents itself as an iterable of outgoing batches.\r\n",
"\r\n",
" Args:\r\n",
" src_iterable: An iterable (e.g., a DataLoader), whose iterator yields one incoming batch,\r\n",
" where a batch is a list of unpadded, variable-length Sequences of token\r\n",
" IDs. Since this only needs to be an Iterable, it could also be a generator object\r\n",
" like the result of `itertools.batched(dataset_list,batch_size))`\r\n",
"\r\n",
" src_batch_size: This is the INCOMING batch size, the number of seqs in one batch yielded\r\n",
" from `src_iterable`'s iterator.\r\n",
"\r\n",
" src_max_seq_len: The maximum number of tokens in a seq within an incoming batch.\r\n",
"\r\n",
" out_batch_size: the number of pseqs (packed seqs) in one outgoing batch\r\n",
"\r\n",
" out_pseq_len: the number of tokens per packed seq, in every outgoing batch\r\n",
"\r\n",
" buffer_size: The maximum number of seqs which may be buffered internally.\r\n",
"\r\n",
" pad_token_id: The token ID used for padding the space which cannot be filled to reach out_pseq_len.\r\n",
"\r\n",
" mask_token_id: The token ID used for masking tokens in the input sequence.\r\n",
"\r\n",
" ignore_token_id: The token ID used to ignore tokens. Expected to be applied to every non-masked token, so the model only trains on predictions of masked tokens.\r\n",
"\r\n",
" suppress_masking: If True, the sequence packer will not perform masked language modeling.\r\n",
"\r\n",
" batch_size_warmup_min_size: If not None, the sequence packer will gradually increase the batch size from batch_size_warmup_min_size to out_batch_size over the course of the warmup_tokens.\r\n",
" batch_size_warmup_min_size must be a multiple of micro_batch_size.\r\n",
"\r\n",
" batch_size_warmup_tokens: If not None, the sequence packer will gradually increase the batch size from batch_size_warmup_min_size to out_batch_size over the course of the warmup_tokens.\r\n",
"\r\n",
" masking_ratio_cooldown_tokens: If not None, the sequence packer will gradually decrease the masking ratio from mask_prob to final_masking_ratio over the course of the cooldown_tokens.\r\n",
"\r\n",
" masking_ratio_cooldown_start_tokens: The number of tokens to wait before starting the masking ratio cooldown.\r\n",
"\r\n",
" final_masking_ratio: The final masking ratio to cooldown to.\r\n",
"\r\n",
" masking_cooldown_schedule_type: The type of cooldown schedule to use for the masking ratio. Can be 'linear' or 'stair'.\r\n",
"\r\n",
" masking_ratio_step_size: The step size for the `stair` cooldown schedule, in percentage terms (e.g., 0.05 for 5%).\r\n",
"\r\n",
" world_size: The number of processes participating in this training run. batch_size_warmup_min_size is divided by this number.\r\n",
" \"\"\"\r\n",
" assert buffer_size >= out_batch_size, f\"required that {buffer_size=} >= {out_batch_size=}\"\r\n",
" self.src_dataloader_len = len(src_iterable)\r\n",
" self.src_iterable = src_iterable\r\n",
" self.src_batch_size = src_batch_size\r\n",
" self.out_batch_size = out_batch_size\r\n",
" self.out_pseq_len = out_pseq_len\r\n",
" self.buffer_size = buffer_size\r\n",
" self.pad_token_id = pad_token_id\r\n",
" self.mask_token_id = mask_token_id\r\n",
" self.ignore_token_id = ignore_token_id\r\n",
" self.mask_prob = mask_prob\r\n",
" self.suppress_masking = suppress_masking\r\n",
" # internals\r\n",
" self.buffer = deque() # internal buffer holds individual seqs, as tensors.\r\n",
" # for stats to report packing efficiency.\r\n",
" self._seqs_consumed = 0\r\n",
" self._seqs_emitted = 0\r\n",
" # Set random seed\r\n",
" self.seed = seed\r\n",
" self.epoch = -1\r\n",
" self._token_count = 0\r\n",
" self.batch_size_scheduler = None\r\n",
" if batch_size_warmup_min_size is not None and batch_size_warmup_tokens is not None:\r\n",
" self.batch_size_scheduler = BatchSizeWarmupScheduler(\r\n",
" batch_size_warmup_min_size, out_batch_size, batch_size_warmup_tokens, world_size\r\n",
" )\r\n",
" else:\r\n",
" self.batch_size_scheduler = None\r\n",
" self.masking_ratio_scheduler = None\r\n",
" if masking_ratio_cooldown_tokens is not None and final_masking_ratio is not None:\r\n",
" self.masking_ratio_scheduler = MaskingRatioCooldownScheduler(\r\n",
" initial_masking_ratio=self.mask_prob,\r\n",
" final_masking_ratio=final_masking_ratio,\r\n",
" cooldown_tokens=masking_ratio_cooldown_tokens,\r\n",
" cooldown_start_tokens=masking_ratio_cooldown_start_tokens or 0,\r\n",
" schedule_type=MaskingSchedule(masking_cooldown_schedule_type),\r\n",
" masking_ratio_step_size=masking_ratio_step_size,\r\n",
" world_size=world_size,\r\n",
" )\r\n",
"\r\n",
" @property\r\n",
" def seqs_emitted(self):\r\n",
" \"Number of seqs, incoming from src_iterable, which have been emitted in OUTGOING batches.\"\r\n",
" return self._seqs_emitted\r\n",
"\r\n",
" @property\r\n",
" def seqs_consumed(self):\r\n",
" \"Number of seqs, incoming from src_iterable, which have been consumed.\"\r\n",
" return self._seqs_consumed\r\n",
"\r\n",
" def _reset_state(self):\r\n",
" self.epoch += 1\r\n",
" self.buffer.clear()\r\n",
" self._seqs_consumed = 0\r\n",
" self._seqs_emitted = 0\r\n",
" self.np_rng = np.random.default_rng(self.epoch + self.seed)\r\n",
"\r\n",
" # Update the epoch for the sampler\r\n",
" if isinstance(self.src_iterable, torch.utils.data.dataloader.DataLoader):\r\n",
" if isinstance(self.src_iterable.sampler, torch.utils.data.distributed.DistributedSampler):\r\n",
" self.src_iterable.sampler.set_epoch(self.epoch)\r\n",
"\r\n",
" def __iter__(self):\r\n",
" self._reset_state()\r\n",
" self.src_iterator = iter(self.src_iterable)\r\n",
" return self._generate_batches()\r\n",
"\r\n",
" def __len__(self):\r\n",
" # rather than estimate the packed length of the dataset, we rely on Composer's ability\r\n",
" # to schedule training the using the number of batches or tokens instead of epochs.\r\n",
" return None\r\n",
"\r\n",
" def _fill_buffer(self, max_items_to_add=float(\"inf\")) -> int:\r\n",
" \"\"\"\r\n",
" Refills the internal buffer.\r\n",
"\r\n",
" - max_items_to_add: an amount less than or equal to the number of items to add\r\n",
"\r\n",
" Returns: the number of items actually added.\r\n",
"\r\n",
" The default implementation of this simply extends to src.buffer, which is\r\n",
" initialized as a list in __init__. Subclasses which want to use a different data\r\n",
" structure for internal buffering should override this method and also add\r\n",
" code in __init__ to initialize src.buffer appropriately.\r\n",
"\r\n",
" Any implementation of this MUST never place more than self.buffer_size items\r\n",
" in the internal buffer.\r\n",
" \"\"\"\r\n",
" items_added = 0\r\n",
" # NOTE: this should be >=, kept as is to match model training code\r\n",
" # TODO: change if training a new model\r\n",
" while (self.buffer_size - len(self.buffer)) > self.src_batch_size:\r\n",
" try:\r\n",
" # if pulling another batch would fetch more than the requested max, stop\r\n",
" if max_items_to_add < float(\"inf\"):\r\n",
" if (items_added + self.src_batch_size) > max_items_to_add:\r\n",
" # print(\"Not adding, because of max_items_to_fetch\")\r\n",
" break\r\n",
" incoming_batch = next(self.src_iterator)\r\n",
" assert len(incoming_batch) <= self.src_batch_size, (\r\n",
" f\"expected {len(incoming_batch)=} <= {self.src_batch_size=}\"\r\n",
" )\r\n",
" for item in incoming_batch:\r\n",
" if len(item[\"input_ids\"]) > 0: # ignore empty sequences\r\n",
" self.buffer.append(item[\"input_ids\"])\r\n",
" items_added += 1\r\n",
" self._seqs_consumed += 1\r\n",
" except StopIteration:\r\n",
" break\r\n",
" return items_added\r\n",
"\r\n",
" def _generate_batches(self):\r\n",
" \"\"\"\r\n",
" Generates batches of packed sequences.\r\n",
"\r\n",
" The returned generator's iterator will always, when next() is called on it, either:\r\n",
" - return a valid tuple batch (masked_batch, labels, cu_seq_lens,max_seq_lens)\r\n",
" - raise StopIteration\r\n",
" \"\"\"\r\n",
" while True:\r\n",
" retval = self._create_batch()\r\n",
" if retval is None:\r\n",
" break\r\n",
" batch, lst_cu_seq_lens = retval\r\n",
"\r\n",
" assert isinstance(retval, tuple), f\"Unexpected {type(retval)=}\"\r\n",
" assert isinstance(retval[0], np.ndarray), f\"Unexpected {type(retval[0])=}\"\r\n",
" assert isinstance(retval[1], list), f\"Unexpected {type(retval[1])=}\"\r\n",
"\r\n",
" cu_seq_lens = [torch.tensor(x, dtype=torch.int32) for x in lst_cu_seq_lens]\r\n",
" max_seq_lens = [torch.max(x[1:] - x[:-1]).item() for x in cu_seq_lens]\r\n",
" assert isinstance(cu_seq_lens, list), f\"Unexpected {type(cu_seq_lens)=}\"\r\n",
" if self.suppress_masking:\r\n",
" yieldval = {\r\n",
" \"input_ids\": torch.from_numpy(batch),\r\n",
" \"labels\": None,\r\n",
" \"cu_seqlens\": cu_seq_lens,\r\n",
" \"max_seqlen\": max_seq_lens,\r\n",
" }\r\n",
" else:\r\n",
" mask_prob = self.mask_prob\r\n",
" if self.masking_ratio_scheduler:\r\n",
" mask_prob = self.masking_ratio_scheduler(self._token_count)\r\n",
" (masked_batch, labels) = SequencePacker.mlm_masking(\r\n",
" batch, mask_prob, self.mask_token_id, self.pad_token_id, self.ignore_token_id, self.np_rng\r\n",
" )\r\n",
" yieldval = {\r\n",
" \"input_ids\": torch.from_numpy(masked_batch),\r\n",
" \"labels\": torch.from_numpy(labels),\r\n",
" \"cu_seqlens\": cu_seq_lens,\r\n",
" \"max_seqlen\": max_seq_lens,\r\n",
" \"attention_mask\": torch.from_numpy(np.where(batch == self.pad_token_id, 0, 1)),\r\n",
" }\r\n",
" if self.masking_ratio_scheduler:\r\n",
" yieldval[\"mask_prob\"] = torch.tensor(mask_prob, dtype=torch.float32)\r\n",
" self._token_count += yieldval[\"attention_mask\"].sum().item()\r\n",
" # # assert isinstance(yieldval[0], torch.Tensor), f\"Unexpected {type(yieldval[0])=}\"\r\n",
" # if not self.suppress_masking:\r\n",
" # assert isinstance(yieldval[1], torch.Tensor), f\"Unexpected {type(yieldval[1])=}\"\r\n",
" # assert isinstance(yieldval[2], list), f\"Unexpected {type(yieldval[2])=}\"\r\n",
" # if yieldval[2]:\r\n",
" # assert isinstance(yieldval[2][0], torch.Tensor), f\"Unexpected {type(yieldval[2][0])=}\"\r\n",
" yield yieldval\r\n",
"\r\n",
" @staticmethod\r\n",
" def mlm_masking(\r\n",
" seq: np.ndarray,\r\n",
" mask_prob: float,\r\n",
" mask_token: int,\r\n",
" pad_token: int = -1,\r\n",
" ignore_index: int = -100,\r\n",
" np_rng=np.random.default_rng(),\r\n",
" ) -> tuple[np.ndarray, np.ndarray]:\r\n",
" \"\"\"\r\n",
" Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.\r\n",
"\r\n",
" This is exactly a numpy version of transformers' `DataCollatorForLanguageModeling.torch_mask_tokens`\r\n",
" https://github.com/huggingface/transformers/blob/main/src/transformers/data/data_collator.py#L827\r\n",
"\r\n",
" It performs masking in a way that produces on expectation the following masked inputs:\r\n",
" - (1-mask_prob) of the original positions will be untouched.\r\n",
" - mask_prob * 80% of the original positions get replaced with a mask token\r\n",
" - mask_prob * 10% of the original positions get replaced with a random token\r\n",
" - mask_prob * 10% of the original positions also remain untouched.\r\n",
" This generates the masked_inputs.\r\n",
"\r\n",
" It also generates a labels array, which has ignore tokens in the (1-mask_prob) positions\r\n",
"\r\n",
" These proportions are expectation values since the random transformation is performed\r\n",
" independently per element. (This is why it is agnostic wrt shape.)\r\n",
"\r\n",
" Args:\r\n",
" seq (np.ndarray): the input token IDs (e.g., a sequence, or batch of seqs)\r\n",
" mask_prob (float): probability of initially masking a token, in the first \"wave\" of masking\r\n",
" mask_token (int): token to use for masking\r\n",
" ignore_index (int): the token indicating that position should be ignored during training. We call it `ignore_index` to conform to the API of the cross entropy loss function.\r\n",
"\r\n",
" Returns:\r\n",
" tuple[np.array,np.array]: (masked_seq, labels)\r\n",
" masked_seq: the input seq with some tokens replaced by `mask_token`\r\n",
" labels: the original input seq with non-masked tokens replaced by `ignore_index`\r\n",
" \"\"\"\r\n",
" # Create labels\r\n",
" labels = np.where(seq == pad_token, ignore_index, seq)\r\n",
"\r\n",
" # Create a single mask\r\n",
" rand = np_rng.random(seq.shape)\r\n",
"\r\n",
" # Partition the probability space appropriately using a single mask\r\n",
" # 80% of the time, we mask the token\r\n",
" mask_mask = rand < mask_prob * 0.8\r\n",
" # 10% of the time, we replace the token with a random token\r\n",
" random_mask = (rand >= mask_prob * 0.8) & (rand < mask_prob * 0.9)\r\n",
" # 10% of the time, we keep the token the same\r\n",
" keep_mask = (rand >= mask_prob * 0.9) & (rand < mask_prob)\r\n",
"\r\n",
" # We only compute loss over the tokens marked for masking\r\n",
" labels = np.where(mask_mask | random_mask | keep_mask, labels, ignore_index)\r\n",
"\r\n",
" # Apply masking\r\n",
" seq = np.where(mask_mask, mask_token, seq)\r\n",
"\r\n",
" # Apply random replacement\r\n",
" random_words = np_rng.integers(0, np.max(seq) + 1, size=seq.shape)\r\n",
" seq = np.where(random_mask, random_words, seq)\r\n",
"\r\n",
" return seq, labels\r\n",
"\r\n",
" @abstractmethod\r\n",
" def _create_batch(self) -> Optional[tuple[np.ndarray, list[list[int]]]]:\r\n",
" \"\"\"\r\n",
" Returns a batch of packed sequences with its cumulative seq length information.\r\n",
"\r\n",
" Or else, returns None if it cannot build a full outgoing batch.\r\n",
"\r\n",
" Must mutate self.buffer to remove the sequences that are packed into the batch.\r\n",
"\r\n",
" Returns:\r\n",
" (out_batch,cumulative_seq_len):tuple[torch.tensor, list[list[int]]]\r\n",
" where:\r\n",
" - out_batch is a tensor of shape (out_batch_size, out_pseq_len);\r\n",
" - cum_seq_lens is a list of lists, where the outer list is of len out_batch_size,\r\n",
" and each inner list is of varying length, and contains the start positions of\r\n",
" every seq in the pseq, and the end position of the last seq in the pseq. This end\r\n",
" position is necessary to communicate if any padding tokens were added.\r\n",
" \"\"\"\r\n",
" pass\r\n",
"\r\n",
"\r\n",
"@njit\r\n",
"def find_best_fit(remaining_spaces, seq_len):\r\n",
" valid_spaces = seq_len <= remaining_spaces\r\n",
" if np.any(valid_spaces):\r\n",
" valid_space_sizes = remaining_spaces[valid_spaces]\r\n",
" best_fit_idx = np.argmin(valid_space_sizes)\r\n",
" return np.arange(len(remaining_spaces))[valid_spaces][best_fit_idx]\r\n",
" return -1\r\n",
"\r\n",
"\r\n",
"class GreedyBestFitSequencePacker(SequencePacker):\r\n",
" @classmethod\r\n",
" def from_composer(\r\n",
" cls,\r\n",
" src_iterable: Iterable[list[list[int]]],\r\n",
" batch_size: int = 512,\r\n",
" micro_batch_size: int = 32,\r\n",
" max_seq_len: int = 1024,\r\n",
" buffer_size: int = 5120,\r\n",
" # token values\r\n",
" pad_token_id: int = -1,\r\n",
" mask_token_id: int = 0,\r\n",
" ignore_token_id: int = -100,\r\n",
" mask_prob: float = 0.3,\r\n",
" # transform values\r\n",
" seed=42,\r\n",
" suppress_masking=False,\r\n",
" batch_size_warmup_min_size: Optional[int] = None,\r\n",
" batch_size_warmup_tokens: Optional[Union[str, Time]] = None,\r\n",
" masking_ratio_cooldown_tokens: Optional[Union[str, Time]] = None,\r\n",
" masking_ratio_cooldown_start_tokens: Optional[Union[str, Time]] = None,\r\n",
" final_masking_ratio: Optional[float] = None,\r\n",
" masking_cooldown_schedule_type: Union[str, MaskingSchedule] = \"linear\",\r\n",
" masking_ratio_step_size: Optional[float] = None,\r\n",
" world_size: int = 1,\r\n",
" ) -> \"GreedyBestFitSequencePacker\":\r\n",
" if batch_size_warmup_min_size is not None:\r\n",
" if batch_size_warmup_min_size % micro_batch_size != 0:\r\n",
" raise ValueError(f\"{batch_size_warmup_min_size=} must be a multiple of {micro_batch_size=}\")\r\n",
" batch_size_warmup_min_size = int(batch_size_warmup_min_size / micro_batch_size)\r\n",
" return cls(\r\n",
" # input shape\r\n",
" src_iterable=src_iterable,\r\n",
" src_batch_size=batch_size,\r\n",
" src_max_seq_len=max_seq_len,\r\n",
" # output shape\r\n",
" out_batch_size=int(batch_size / micro_batch_size),\r\n",
" out_pseq_len=int(micro_batch_size * max_seq_len),\r\n",
" # internal\r\n",
" buffer_size=buffer_size,\r\n",
" # transformation\r\n",
" pad_token_id=pad_token_id,\r\n",
" mask_token_id=mask_token_id,\r\n",
" ignore_token_id=ignore_token_id,\r\n",
" mask_prob=mask_prob,\r\n",
" seed=seed,\r\n",
" suppress_masking=suppress_masking,\r\n",
" batch_size_warmup_min_size=batch_size_warmup_min_size,\r\n",
" batch_size_warmup_tokens=batch_size_warmup_tokens,\r\n",
" masking_ratio_cooldown_tokens=masking_ratio_cooldown_tokens,\r\n",
" masking_ratio_cooldown_start_tokens=masking_ratio_cooldown_start_tokens,\r\n",
" final_masking_ratio=final_masking_ratio,\r\n",
" masking_cooldown_schedule_type=masking_cooldown_schedule_type,\r\n",
" masking_ratio_step_size=masking_ratio_step_size,\r\n",
" world_size=world_size,\r\n",
" )\r\n",
"\r\n",
" def _create_batch(self) -> Optional[tuple[np.ndarray, list[list[int]]]]:\r\n",
" if self.batch_size_scheduler:\r\n",
" self.out_batch_size = self.batch_size_scheduler(self._token_count)\r\n",
"\r\n",
" batch = np.full(\r\n",
" (self.out_batch_size, self.out_pseq_len), self.pad_token_id, dtype=np.int64\r\n",
" ) # the pseqs being constructed\r\n",
" seq_counts = np.zeros(self.out_batch_size, dtype=np.int32) # the count of seqs per pseq\r\n",
" cum_seq_lens = [[0] for _ in range(self.out_batch_size)]\r\n",
" remaining_spaces = np.full(\r\n",
" (self.out_batch_size,), self.out_pseq_len, dtype=np.int32\r\n",
" ) # the space remaining per pseq\r\n",
" temp_buffer = []\r\n",
"\r\n",
" while True:\r\n",
" # Check if buffer has more items, and if not replenish\r\n",
" if not self.buffer:\r\n",
" items_to_fetch = self.buffer_size - len(temp_buffer)\r\n",
" items_added = self._fill_buffer(items_to_fetch)\r\n",
" if items_added == 0:\r\n",
" break\r\n",
"\r\n",
" seq = self.buffer.popleft()\r\n",
" seq_len = len(seq)\r\n",
"\r\n",
" # Find the best fit (smallest space that can accommodate the sequence)\r\n",
" best_fit_idx = find_best_fit(remaining_spaces, seq_len)\r\n",
" if best_fit_idx != -1:\r\n",
" end_pos = self.out_pseq_len - remaining_spaces[best_fit_idx]\r\n",
" batch[best_fit_idx, end_pos : end_pos + seq_len] = seq\r\n",
" seq_counts[best_fit_idx] += 1\r\n",
" remaining_spaces[best_fit_idx] -= seq_len\r\n",
" cum_seq_lens[best_fit_idx].append(cum_seq_lens[best_fit_idx][-1] + seq_len)\r\n",
" else:\r\n",
" # Can't fit the sequence, save for next batch\r\n",
" temp_buffer.append(seq)\r\n",
"\r\n",
" # Add any sequences we skipped back to the start of the buffer\r\n",
" self.buffer.extendleft(temp_buffer)\r\n",
"\r\n",
" if np.all(seq_counts > 0):\r\n",
" self._seqs_emitted += np.sum(seq_counts)\r\n",
" for x in cum_seq_lens:\r\n",
" if x[-1] != self.out_pseq_len:\r\n",
" x.append(self.out_pseq_len)\r\n",
" return batch, cum_seq_lens\r\n",
" else:\r\n",
" # If we can't form a full batch, we return None to signal the end\r\n",
" return None\r\n",
"\r\n",
"\r\n",
"T = TypeVar(\"T\")\r\n",
"\r\n",
"\r\n",
"class BufferedIterable(Generic[T]):\r\n",
" def __init__(self, iterable: Iterable[T], buffer_size: int):\r\n",
" \"\"\"\r\n",
" Args:\r\n",
" - iterable: an object which generates a fresh iterator on iter() and which implements len()\r\n",
" \"\"\"\r\n",
" self.iterable = iterable\r\n",
" self.buffer_size = buffer_size\r\n",
"\r\n",
" def __iter__(self):\r\n",
" return BufferedIterator(self.iterable, self.buffer_size)\r\n",
"\r\n",
" @property\r\n",
" def batch_size(self) -> int:\r\n",
" if isinstance(self.iterable, SequencePacker):\r\n",
" return self.iterable.out_batch_size\r\n",
" else:\r\n",
" raise TypeError(\"Expected a SequencePacker\")\r\n",
"\r\n",
"\r\n",
"class BufferedIterator(Generic[T]):\r\n",
" def __init__(self, iterable: Iterable[T], buffer_size: int):\r\n",
" self.iterator = iter(iterable)\r\n",
" self.buffer = deque(maxlen=buffer_size)\r\n",
" self.buffer_size = buffer_size\r\n",
" self.lock = threading.Lock()\r\n",
" self.exhausted = False\r\n",
" self.filler_thread = threading.Thread(target=self._background_fill, daemon=True)\r\n",
" self.filler_thread.start()\r\n",
"\r\n",
" def _background_fill(self):\r\n",
" # Fill up the buffer, whenever possible, in the background\r\n",
" while not self.exhausted:\r\n",
" if len(self.buffer) < self.buffer_size:\r\n",
" try:\r\n",
" item = next(self.iterator)\r\n",
" with self.lock:\r\n",
" self.buffer.append(item)\r\n",
" except StopIteration:\r\n",
" self.exhausted = True\r\n",
" break\r\n",
" else:\r\n",
" time.sleep(0.01) # Sleep for a bit to avoid busy waiting\r\n",
"\r\n",
" def __iter__(self):\r\n",
" return self\r\n",
"\r\n",
" def __next__(self) -> T:\r\n",
" while True:\r\n",
" if not self.buffer:\r\n",
" if self.exhausted:\r\n",
" # We've exhausted the iterator and the buffer so we're done\r\n",
" raise StopIteration\r\n",
" else:\r\n",
" # The buffer is empty but the iterator is not exhausted yet.\r\n",
" # Let's give the filler thread a chance to add items to the buffer\r\n",
" time.sleep(0.01)\r\n",
" else:\r\n",
" with self.lock:\r\n",
" return self.buffer.popleft()\r\n",
"\r\n",
"\r\n",
"def split_packed_batch(\r\n",
" batch: Batch, microbatch_size: Union[int, float], padding_tolerance: float = 1.0, mark_dynamic: bool = True\r\n",
") -> Sequence:\r\n",
" # NOTE: Packed sequences are already packed into a microbatch size worth of tokens.\r\n",
" # So to correctly return a microbatch worth of data, we will simply return each item (i.e. microbatch_size 1)\r\n",
"\r\n",
" num_items = batch[\"input_ids\"].shape[0]\r\n",
" split_inputs = [x.squeeze() for x in batch[\"input_ids\"].split(1)]\r\n",
" split_labels = [x.squeeze() for x in batch[\"labels\"].split(1)]\r\n",
" split_attention_masks = [x.squeeze() for x in batch[\"attention_mask\"].split(1)]\r\n",
" split_cu_seqlens = batch[\"cu_seqlens\"]\r\n",
" result = []\r\n",
" for i in range(num_items):\r\n",
" attention_mask = split_attention_masks[i]\r\n",
" padding_amount = 1 - (attention_mask.sum() / len(attention_mask))\r\n",
"\r\n",
" if padding_amount > padding_tolerance:\r\n",
" last_non_pad = attention_mask.nonzero().max()\r\n",
" input_ids = split_inputs[i][: last_non_pad + 1]\r\n",
" labels = split_labels[i][: last_non_pad + 1]\r\n",
" cu_seqlens = split_cu_seqlens[i][:-1]\r\n",
" attention_mask = attention_mask[: last_non_pad + 1]\r\n",
" else:\r\n",
" input_ids = split_inputs[i]\r\n",
" labels = split_labels[i]\r\n",
" cu_seqlens = split_cu_seqlens[i]\r\n",
" if mark_dynamic:\r\n",
" torch._dynamo.mark_dynamic(cu_seqlens, index=0)\r\n",
" microbatch = {\r\n",
" \"input_ids\": input_ids,\r\n",
" \"labels\": labels,\r\n",
" \"cu_seqlens\": cu_seqlens,\r\n",
" \"max_seqlen\": batch[\"max_seqlen\"][i],\r\n",
" \"attention_mask\": attention_mask,\r\n",
" }\r\n",
" if \"mask_prob\" in batch:\r\n",
" microbatch[\"mask_prob\"] = batch[\"mask_prob\"].clone()\r\n",
" result.append(microbatch)\r\n",
"\r\n",
" assert all([x[\"input_ids\"].shape[-1] == y[\"cu_seqlens\"][-1] for x, y in zip(result, result)])\r\n",
" return result\r\n",
"\r\n",
"\r\n",
"def get_num_tokens_in_packed_batch(batch: Batch, ignore_index: int = -100) -> int:\r\n",
" labels: torch.Tensor | list[torch.Tensor] = batch[\"labels\"]\r\n",
" if isinstance(labels, torch.Tensor):\r\n",
" return (labels != ignore_index).sum().item()\r\n",
" elif isinstance(labels, list):\r\n",
" return sum([(x != ignore_index).sum().item() for x in labels])\r\n",
" else:\r\n",
" raise TypeError('Expected a batch with a \"labels\" key of type list[Tensor] or Tensor')\r\n",
"\r\n",
"\r\n",
"def get_num_samples_in_packed_batch(batch: Batch) -> int:\r\n",
" # Number of sequences can be inferred from cu_seqlens arrays\r\n",
" cu_seqlens: torch.Tensor | list[torch.Tensor] = batch[\"cu_seqlens\"]\r\n",
" if isinstance(cu_seqlens, torch.Tensor):\r\n",
" return cu_seqlens.size()[0] - 1\r\n",
" elif isinstance(cu_seqlens, list):\r\n",
" return sum([x.size()[0] - 1 for x in batch[\"cu_seqlens\"]])\r\n",
" else:\r\n",
" raise TypeError('Expected a batch with a \"cu_seqlens\" key of type list[Tensor] or Tensor')"
]
},
{
"cell_type": "markdown",
"id": "7bebc1d0",
"metadata": {
"heading_collapsed": true,
"input_tokens": 4
},
"source": [
"## Brief overlook"
]
},
{
"cell_type": "markdown",
"id": "94983150",
"metadata": {
"hidden": 1,
"input_tokens": 46,
"output_tokens": 391,
"solveit_ai": true
},
"source": [
"I want to understand how Modernbert sequence packer works. Before I dive into details, Can you give me a brief summary of how it works? \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! Here's a brief overview of how the ModernBERT sequence packer works:\n",
"\n",
"**Main Purpose**: It takes variable-length text sequences and efficiently packs them into fixed-size batches for training, while also applying masked language modeling.\n",
"\n",
"**Key Steps**:\n",
"1. **Input**: Receives batches of variable-length token sequences (like sentences of different lengths)\n",
"2. **Buffering**: Stores sequences in an internal buffer \n",
"3. **Packing**: Uses a \"greedy best-fit\" algorithm to pack multiple short sequences into fixed-length \"packed sequences\" (pseqs)\n",
"4. **Masking**: Applies BERT-style masking (80% mask token, 10% random, 10% unchanged) for training\n",
"5. **Output**: Returns batches of packed sequences with attention information\n",
"\n",
"**Why Pack?**: Instead of padding each sequence individually (wasteful), it combines multiple sequences into one packed sequence, making training much more efficient.\n",
"\n",
"The core insight is that instead of having one sequence per training example, you can have multiple sequences concatenated together, as long as you track where each sequence starts and ends.\n",
"\n",
"Would you like me to explain any of these steps in more detail, or would you prefer to see a simple example of how the packing works?"
]
},
{
"cell_type": "markdown",
"id": "a59311fd",
"metadata": {
"hidden": 1,
"input_tokens": 42,
"output_tokens": 288,
"solveit_ai": true
},
"source": [
"Why does it return the batches of packed sequences with attention information? To use with Flash attention 2, there is no notion of batches.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Excellent observation! You're absolutely right that Flash Attention 2 works differently from standard attention.\n",
"\n",
"The sequence packer returns attention information (specifically `cu_seqlens` and `max_seqlen`) because it's designed to work with **variable-length attention** implementations like Flash Attention 2, which use a different format than traditional batched attention.\n",
"\n",
"Instead of a 2D attention mask, Flash Attention 2 uses:\n",
"- `cu_seqlens`: Cumulative sequence lengths that tell it where each sequence starts/ends within the packed sequence\n",
"- `max_seqlen`: The maximum sequence length in the batch for optimization\n",
"\n",
"This allows Flash Attention 2 to process multiple sequences that are concatenated together without needing traditional padding or 2D attention masks.\n",
"\n",
"Would you like me to show you what these `cu_seqlens` look like with a concrete example, or are you more interested in understanding how the packing algorithm decides which sequences to combine together?"
]
},
{
"cell_type": "markdown",
"id": "c75bce7a",
"metadata": {
"hidden": 1,
"input_tokens": 36,
"output_tokens": 261,
"solveit_ai": true
},
"source": [
"Sequence packer has `out_batch_size` and `out_pseq_len` as arguments. How does it work?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! These two parameters define the shape of the output batches:\n",
"\n",
"- `out_batch_size`: How many packed sequences (pseqs) you want in each output batch\n",
"- `out_pseq_len`: How many tokens each packed sequence should contain\n",
"\n",
"So if you set `out_batch_size=4` and `out_pseq_len=512`, you'll get batches shaped like `(4, 512)` - 4 packed sequences, each containing 512 tokens.\n",
"\n",
"The packer tries to fill each of those 512-token slots by combining multiple shorter sequences. For example, it might pack a 200-token sequence + a 300-token sequence + padding into one 512-token slot.\n",
"\n",
"Would you like me to show you a concrete example of how sequences get packed into these slots, or are you more curious about how it decides which sequences to combine together?"
]
},
{
"cell_type": "markdown",
"id": "2bca73c0",
"metadata": {
"hidden": 1,
"input_tokens": 45,
"output_tokens": 258,
"solveit_ai": true
},
"source": [
"What is the point of returning a batch of packed sequences instead of just one packed sequence? Because Flash attention 2 cannot able to use batch input.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"That's a really insightful question! You're right that Flash Attention 2 processes sequences differently than traditional batched attention.\n",
"\n",
"The reason it returns multiple packed sequences is for **parallelization and memory efficiency**. Even though Flash Attention 2 doesn't use traditional batching internally, the training framework still needs to:\n",
"\n",
"1. **Parallelize across GPUs** - Each GPU can work on different packed sequences simultaneously\n",
"2. **Pipeline processing** - While one packed sequence is being processed, others can be prepared\n",
"3. **Memory management** - Smaller packed sequences are easier to fit in GPU memory than one giant sequence\n",
"\n",
"Think of it like having multiple workers (GPUs) each handling their own \"chunk\" of work, rather than one worker trying to handle everything at once.\n",
"\n",
"Does this make sense, or would you like me to clarify how this relates to the actual training process?"
]
},
{
"cell_type": "markdown",
"id": "9c7b5baa",
"metadata": {
"hidden": 1,
"input_tokens": 99,
"output_tokens": 924,
"solveit_ai": true
},
"source": [
"When using flash attention varlen function, how does performance vary depending on number of sequences and max_seq_len? Let's say we have multiple sequences with total length of 8192 that consist of smaller sequences. Is it more performant to pack smaller sequences together or large and small together? Or it doesn't matter? Or you don't know?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Based on the research I found, I can give you a good answer about Flash Attention varlen performance with different sequence distributions:\n",
"\n",
"Actual throughput and memory improvement depends on the model and the distribution of example lengths in the training data. [^1] [^2] Here's what the research shows:\n",
"\n",
"**Performance varies significantly based on sequence length distribution:**\n",
"\n",
"1. **High variance = better performance**: FLAN has short sequences on average but a large variance in sequence length, so example lengths in each batch may vary widely. This means that padded FLAN batches may incur a significant overhead in unused padding tokens. Training on the FLAN dataset shows a significant benefit using the new DataCollatorWithFlattening in terms of increased throughput. We see a 2x throughput increase on the models shown here: llama2-7B, mistral-7B, and granite-8B-code. [^3]\n",
"\n",
"2. **Low variance = smaller gains**: OrcaMath has longer examples and a lower variance in example length. As such, the improvement from packing is lower. Our experiments show a 1.4x increase in throughput when training using this form of packing on the OrcaMath dataset across these three models. [^4]\n",
"\n",
"**For your 8192 token question**: It generally **doesn't matter much** whether you pack small sequences together vs. mixing large and small, as long as you're efficiently using the available space. Training with data that has a wide variation of example lengths will see the greatest benefit, with respect to padding, by using the DataCollatorWithFlattening. [^5]\n",
"\n",
"The key insight is that the performance benefit comes from **reducing wasted padding tokens**, not from the specific arrangement of sequences within the packed batch.\n",
"\n",
"Would you like me to explain more about why sequence length variance matters so much for Flash Attention performance?\n",
"\n",
"[^1]: https://huggingface.co/blog/packing-with-FA2\n",
"\t\"Actual throughput and memory improvement depends on the model and the distribution of example lengths in the training data. \"\n",
"\n",
"[^2]: https://research.ibm.com/blog/hugging-face-training-flash-attention\n",
"\t\"Actual throughput and memory improvement depends on the model and the distribution of example lengths in the training data. \"\n",
"\n",
"[^3]: https://huggingface.co/blog/packing-with-FA2\n",
"\t\"FLAN has short sequences on average but a large variance in sequence length, so example lengths in each batch may vary widely. This means that padded ...\"\n",
"\n",
"[^4]: https://huggingface.co/blog/packing-with-FA2\n",
"\t\"OrcaMath has longer examples and a lower variance in example length. As such, the improvement from packing is lower. Our experiments show a 1.4x incre...\"\n",
"\n",
"[^5]: https://research.ibm.com/blog/hugging-face-training-flash-attention\n",
"\t\"Training with data that has a wide variation of example lengths will see the greatest benefit, with respect to padding, by using the DataCollatorWithF...\""
]
},
{
"cell_type": "markdown",
"id": "9e54d4e8",
"metadata": {
"hidden": 1,
"input_tokens": 27,
"output_tokens": 229,
"solveit_ai": true
},
"source": [
"Because high seuqnece length variance means there are a lot of padding tokens. \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly right! You've got it.\n",
"\n",
"High sequence length variance means when you use traditional padding (making all sequences the same length), you end up with tons of wasted padding tokens. For example:\n",
"- One sequence: 50 tokens\n",
"- Another sequence: 800 tokens \n",
"- If you pad both to 800, the first sequence has 750 wasted padding tokens!\n",
"\n",
"With packing, you can combine multiple short sequences to fill up the available space more efficiently, dramatically reducing those wasted computations.\n",
"\n",
"That's why datasets like FLAN (with high variance) see 2x speedups, while datasets like OrcaMath (with more uniform lengths) only see 1.4x speedups.\n",
"\n",
"Does this help clarify why the variance matters more than the specific arrangement of sequences?"
]
},
{
"cell_type": "markdown",
"id": "70e73ebd",
"metadata": {
"hidden": 1,
"input_tokens": 54,
"output_tokens": 588,
"solveit_ai": true
},
"source": [
"Tell me about `DataCollatorWithFlattening`. The reference used it with SFT. Why only use SFT? Can it also be used for pretraining, too?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! Based on what I found, `DataCollatorWithFlattening` **can absolutely be used for pretraining too**, not just SFT.\n",
"\n",
"The references focused on SFT because that's where the most immediate benefits are seen, but the technique works for any training that uses Flash Attention 2. By selecting DataCollatorWithFlattening, Hugging Face Trainer users can now seamlessly concatenate sequences into a single tensor while accounting for sequence boundaries during Flash Attention 2 computations. Training with packed instruction tuning examples (without padding) is now compatible with Flash Attention 2 in Hugging Face, thanks to a recent PR and the new DataCollatorWithFlattening. [^1] [^2]\n",
"\n",
"The key insight is that Training with data that has a wide variation of example lengths will see the greatest benefit, with respect to padding, by using the DataCollatorWithFlattening. [^3]\n",
"\n",
"SFT datasets often have high sequence length variance (short prompts + long completions), which is why they highlighted it there. But pretraining datasets can also benefit if they have varied sequence lengths.\n",
"\n",
"Would you like me to explain more about when pretraining would benefit most from this approach?\n",
"\n",
"[^1]: https://research.ibm.com/blog/hugging-face-training-flash-attention\n",
"\t\"By selecting DataCollatorWithFlattening, Hugging Face Trainer users can now seamlessly concatenate sequences into a single tensor while accounting for...\"\n",
"\n",
"[^2]: https://research.ibm.com/blog/hugging-face-training-flash-attention\n",
"\t\"Training with packed instruction tuning examples (without padding) is now compatible with Flash Attention 2 in Hugging Face, thanks to a recent PR and...\"\n",
"\n",
"[^3]: https://research.ibm.com/blog/hugging-face-training-flash-attention\n",
"\t\"Training with data that has a wide variation of example lengths will see the greatest benefit, with respect to padding, by using the DataCollatorWithF...\""
]
},
{
"cell_type": "markdown",
"id": "9bb3f1d8",
"metadata": {
"input_tokens": 4
},
"source": [
"## Detailed Look"
]
},
{
"cell_type": "markdown",
"id": "fa710e0b",
"metadata": {
"input_tokens": 7
},
"source": [
"### mlm_masking"
]
},
{
"cell_type": "markdown",
"id": "3219360b",
"metadata": {
"input_tokens": 34,
"output_tokens": 637,
"solveit_ai": true
},
"source": [
"Can you divide this code into small pieces I can try with? I would prefer one function definition per code block.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great idea! Let's break down the sequence packer code into digestible pieces. \n",
"\n",
"Let's start with the simplest utility functions. Here's the first one - the core masking logic:\n",
"\n",
"```python\n",
"@staticmethod\n",
"def mlm_masking(\n",
" seq: np.ndarray,\n",
" mask_prob: float,\n",
" mask_token: int,\n",
" pad_token: int = -1,\n",
" ignore_index: int = -100,\n",
" np_rng=np.random.default_rng(),\n",
") -> tuple[np.ndarray, np.ndarray]:\n",
" \"\"\"\n",
" Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.\n",
" \"\"\"\n",
" # Create labels\n",
" labels = np.where(seq == pad_token, ignore_index, seq)\n",
"\n",
" # Create a single mask\n",
" rand = np_rng.random(seq.shape)\n",
"\n",
" # Partition the probability space appropriately using a single mask\n",
" # 80% of the time, we mask the token\n",
" mask_mask = rand < mask_prob * 0.8\n",
" # 10% of the time, we replace the token with a random token\n",
" random_mask = (rand >= mask_prob * 0.8) & (rand < mask_prob * 0.9)\n",
" # 10% of the time, we keep the token the same\n",
" keep_mask = (rand >= mask_prob * 0.9) & (rand < mask_prob)\n",
"\n",
" # We only compute loss over the tokens marked for masking\n",
" labels = np.where(mask_mask | random_mask | keep_mask, labels, ignore_index)\n",
"\n",
" # Apply masking\n",
" seq = np.where(mask_mask, mask_token, seq)\n",
"\n",
" # Apply random replacement\n",
" random_words = np_rng.integers(0, np.max(seq) + 1, size=seq.shape)\n",
" seq = np.where(random_mask, random_words, seq)\n",
"\n",
" return seq, labels\n",
"```\n",
"\n",
"This function implements the masked language modeling strategy. Would you like to try this first and see how it works, or would you prefer to start with a different piece?"
]
},
{
"cell_type": "markdown",
"id": "05698272",
"metadata": {
"input_tokens": 12,
"output_tokens": 285,
"solveit_ai": true
},
"source": [
"What does `@staticmethod` do?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"`@staticmethod` is a Python decorator that creates a static method. A static method belongs to the class but doesn't need access to `self` (the instance) or `cls` (the class itself).\n",
"\n",
"Here's the key difference:\n",
"\n",
"```python\n",
"class MyClass:\n",
" def regular_method(self): # needs an instance\n",
" return \"I need self\"\n",
" \n",
" @staticmethod\n",
" def static_method(): # doesn't need an instance\n",
" return \"I don't need self\"\n",
"```\n",
"\n",
"You can call a static method directly on the class without creating an instance:\n",
"\n",
"```python\n",
"result = MyClass.static_method() # No instance needed!\n",
"```\n",
"\n",
"In the sequence packer code, `mlm_masking` is static because it's a utility function that just processes arrays - it doesn't need any instance data.\n",
"\n",
"Would you like to see how this works in practice, or shall we move on to understanding what the masking function actually does?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9152b129",
"metadata": {
"input_tokens": 72
},
"outputs": [],
"source": [
"class MyClass:\n",
" def regular_method(self): # needs an instance\n",
" return \"I need self\"\n",
" \n",
" @staticmethod\n",
" def static_method(): # doesn't need an instance\n",
" return \"I don't need self\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "14907eca",
"metadata": {
"input_tokens": 19
},
"outputs": [],
"source": [
"result = MyClass.static_method() # No instance needed!"
]
},
{
"cell_type": "markdown",
"id": "4a20657c",
"metadata": {
"input_tokens": 63,
"output_tokens": 256,
"solveit_ai": true
},
"source": [
"What is the purpose of having this inside of a class rather than a separate function? If it does not use any variable from `self` or the class, is it still good to create a staticmethod?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! You're absolutely right to wonder about this. There are a few reasons why someone might put a static method inside a class:\n",
"\n",
"**Organizational reasons:**\n",
"- It groups related functionality together (the masking function belongs conceptually with the SequencePacker)\n",
"- It creates a namespace (you know `SequencePacker.mlm_masking` is for sequence packing)\n",
"\n",
"**But you're onto something important** - if it doesn't use class data, it could just be a standalone function! In fact, many Python developers would argue that's cleaner.\n",
"\n",
"The function could easily be written as:\n",
"\n",
"```python\n",
"def mlm_masking(seq, mask_prob, mask_token, ...):\n",
" # same implementation\n",
"```\n",
"\n",
"What do you think would be better in this case - keeping it as a static method for organization, or making it a standalone function for simplicity?"
]
},
{
"cell_type": "markdown",
"id": "f41e5c40",
"metadata": {
"input_tokens": 16,
"output_tokens": 109,
"solveit_ai": true
},
"source": [
"I think it is okay to leave as it is.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! That's a reasonable choice - keeping it organized with the class makes sense.\n",
"\n",
"Now, would you like to try out the `mlm_masking` function to see how it works? We could create a simple example with some token sequences to understand what the masking does.\n",
"\n",
"Or would you prefer to look at the next piece of the sequence packer code first?"
]
},
{
"cell_type": "markdown",
"id": "c3ac9d4b",
"metadata": {
"input_tokens": 9,
"output_tokens": 145,
"solveit_ai": true
},
"source": [
"Provide some simple examples plz\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great! Let's create a simple example to see how the masking works. Here's a small test:\n",
"\n",
"```python\n",
"import numpy as np\n",
"\n",
"# Simple example sequence: [1, 2, 3, 4, 5]\n",
"seq = np.array([1, 2, 3, 4, 5])\n",
"```\n",
"\n",
"Try running this first, then we can call the masking function on it. What do you think will happen when we apply masking to this sequence?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9c889434",
"metadata": {
"input_tokens": 66
},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"# Simple example sequence: [1, 2, 3, 4, 5]\n",
"seq = np.array([1, 2, 3, 4, 5])"
]
},
{
"cell_type": "markdown",
"id": "b61e09df",
"metadata": {
"input_tokens": 4,
"output_tokens": 480,
"solveit_ai": true
},
"source": [
"seq becomes masked\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly! Now let's see the masking in action. First we need to add the `mlm_masking` function. Here it is:\n",
"\n",
"```python\n",
"def mlm_masking(\n",
" seq: np.ndarray,\n",
" mask_prob: float,\n",
" mask_token: int,\n",
" pad_token: int = -1,\n",
" ignore_index: int = -100,\n",
" np_rng=np.random.default_rng(),\n",
") -> tuple[np.ndarray, np.ndarray]:\n",
" # Create labels\n",
" labels = np.where(seq == pad_token, ignore_index, seq)\n",
"\n",
" # Create a single mask\n",
" rand = np_rng.random(seq.shape)\n",
"\n",
" # Partition the probability space appropriately using a single mask\n",
" mask_mask = rand < mask_prob * 0.8\n",
" random_mask = (rand >= mask_prob * 0.8) & (rand < mask_prob * 0.9)\n",
" keep_mask = (rand >= mask_prob * 0.9) & (rand < mask_prob)\n",
"\n",
" # We only compute loss over the tokens marked for masking\n",
" labels = np.where(mask_mask | random_mask | keep_mask, labels, ignore_index)\n",
"\n",
" # Apply masking\n",
" seq = np.where(mask_mask, mask_token, seq)\n",
"\n",
" # Apply random replacement\n",
" random_words = np_rng.integers(0, np.max(seq) + 1, size=seq.shape)\n",
" seq = np.where(random_mask, random_words, seq)\n",
"\n",
" return seq, labels\n",
"```\n",
"\n",
"Now try calling it on your sequence. What parameters do you think we should use?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4cec3bf7",
"metadata": {
"input_tokens": 403
},
"outputs": [],
"source": [
"def mlm_masking(\n",
" seq: np.ndarray,\n",
" mask_prob: float,\n",
" mask_token: int,\n",
" pad_token: int = -1,\n",
" ignore_index: int = -100,\n",
" np_rng=np.random.default_rng(),\n",
") -> tuple[np.ndarray, np.ndarray]:\n",
" # Create labels\n",
" labels = np.where(seq == pad_token, ignore_index, seq)\n",
"\n",
" # Create a single mask\n",
" rand = np_rng.random(seq.shape)\n",
"\n",
" # Partition the probability space appropriately using a single mask\n",
" mask_mask = rand < mask_prob * 0.8\n",
" random_mask = (rand >= mask_prob * 0.8) & (rand < mask_prob * 0.9)\n",
" keep_mask = (rand >= mask_prob * 0.9) & (rand < mask_prob)\n",
"\n",
" # We only compute loss over the tokens marked for masking\n",
" labels = np.where(mask_mask | random_mask | keep_mask, labels, ignore_index)\n",
"\n",
" # Apply masking\n",
" seq = np.where(mask_mask, mask_token, seq)\n",
"\n",
" # Apply random replacement\n",
" random_words = np_rng.integers(0, np.max(seq) + 1, size=seq.shape)\n",
" seq = np.where(random_mask, random_words, seq)\n",
"\n",
" return seq, labels"
]
},
{
"cell_type": "markdown",
"id": "9d6bf842",
"metadata": {
"input_tokens": 13,
"output_tokens": 153,
"solveit_ai": true
},
"source": [
"We use seq, mask_prob, mask_token\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! Let's try it with some reasonable values:\n",
"\n",
"```python\n",
"masked_seq, labels = mlm_masking(seq, mask_prob=0.3, mask_token=0)\n",
"print(\"Original:\", seq)\n",
"print(\"Masked: \", masked_seq)\n",
"print(\"Labels: \", labels)\n",
"```\n",
"\n",
"What do you expect to see? Remember, with 30% masking probability, some tokens should be replaced with 0 (the mask token), some might become random numbers, and some stay the same."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a3ae945c",
"metadata": {
"input_tokens": 69,
"output_tokens": 82
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Original: [1 2 3 4 5]\n",
"Masked: [1 2 0 0 5]\n",
"Labels: [-100 -100 3 4 -100]\n"
]
}
],
"source": [
"masked_seq, labels = mlm_masking(seq, mask_prob=0.3, mask_token=0)\n",
"print(\"Original:\", seq)\n",
"print(\"Masked: \", masked_seq)\n",
"print(\"Labels: \", labels)"
]
},
{
"cell_type": "markdown",
"id": "b0f3c3f7",
"metadata": {
"input_tokens": 7
},
"source": [
"#### Line by line details"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0526bd9d",
"metadata": {
"input_tokens": 61,
"output_tokens": 87
},
"outputs": [
{
"data": {
"text/plain": [
"array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, -1, -1])"
]
},
"execution_count": 92,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"seq = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, -1, -1])\r\n",
"seq"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c40dd295",
"metadata": {
"input_tokens": 52
},
"outputs": [],
"source": [
"mask_prob = 0.3\r\n",
"mask_token = 0\r\n",
"pad_token = -1\r\n",
"ignore_index = -100\r\n",
"np_rng=np.random.default_rng(42)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac0f67ae",
"metadata": {
"input_tokens": 22,
"output_tokens": 91
},
"outputs": [
{
"data": {
"text/plain": [
"array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, -100,\n",
" -100])"
]
},
"execution_count": 96,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"labels = np.where(seq == pad_token, ignore_index, seq)\r\n",
"labels"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4607a880",
"metadata": {
"input_tokens": 13,
"output_tokens": 148
},
"outputs": [
{
"data": {
"text/plain": [
"array([0.77395605, 0.43887844, 0.85859792, 0.69736803, 0.09417735,\n",
" 0.97562235, 0.7611397 , 0.78606431, 0.12811363, 0.45038594,\n",
" 0.37079802, 0.92676499])"
]
},
"execution_count": 98,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"rand = np_rng.random(seq.shape)\r\n",
"rand"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "408eb7ad",
"metadata": {
"input_tokens": 22,
"output_tokens": 61
},
"outputs": [
{
"data": {
"text/plain": [
"array([False, False, False, False, True, False, False, False, True,\n",
" False, False, False])"
]
},
"execution_count": 100,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"mask_mask = rand < mask_prob * 0.8\r\n",
"mask_mask"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cda3b188",
"metadata": {
"input_tokens": 42,
"output_tokens": 58
},
"outputs": [
{
"data": {
"text/plain": [
"array([False, False, False, False, False, False, False, False, False,\n",
" False, False, False])"
]
},
"execution_count": 102,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"random_mask = (rand >= mask_prob * 0.8) & (rand < mask_prob * 0.9)\n",
"random_mask"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9a3f7f68",
"metadata": {
"input_tokens": 34,
"output_tokens": 58
},
"outputs": [
{
"data": {
"text/plain": [
"array([False, False, False, False, False, False, False, False, False,\n",
" False, False, False])"
]
},
"execution_count": 104,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"keep_mask = (rand >= mask_prob * 0.9) & (rand < mask_prob)\r\n",
"keep_mask"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "313eaf92",
"metadata": {
"input_tokens": 46,
"output_tokens": 78
},
"outputs": [
{
"data": {
"text/plain": [
"array([-100, -100, -100, -100, 5, -100, -100, -100, 9, -100, -100,\n",
" -100])"
]
},
"execution_count": 106,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# We only compute loss over the tokens marked for masking\r\n",
"labels = np.where(mask_mask | random_mask | keep_mask, labels, ignore_index)\r\n",
"labels"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7666bd1f",
"metadata": {
"input_tokens": 25,
"output_tokens": 87
},
"outputs": [
{
"data": {
"text/plain": [
"array([ 1, 2, 3, 4, 0, 6, 7, 8, 0, 10, -1, -1])"
]
},
"execution_count": 108,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Apply masking\r\n",
"seq = np.where(mask_mask, mask_token, seq)\r\n",
"seq"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cadf11d8",
"metadata": {
"input_tokens": 46,
"output_tokens": 73
},
"outputs": [
{
"data": {
"text/plain": [
"array([8, 7, 4, 9, 5, 4, 4, 2, 1, 6, 9, 0])"
]
},
"execution_count": 110,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Apply random replacement\r\n",
"random_words = np_rng.integers(0, np.max(seq) + 1, size=seq.shape)\r\n",
"random_words"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ecd41d0b",
"metadata": {
"input_tokens": 19,
"output_tokens": 87
},
"outputs": [
{
"data": {
"text/plain": [
"array([ 1, 2, 3, 4, 0, 6, 7, 8, 0, 10, -1, -1])"
]
},
"execution_count": 112,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"seq = np.where(random_mask, random_words, seq)\r\n",
"seq"
]
},
{
"cell_type": "markdown",
"id": "36fe81c7",
"metadata": {
"input_tokens": 39,
"output_tokens": 330,
"solveit_ai": true
},
"source": [
"I don't understand what is the purpose of having `mask_mask`, `random_mask`, and `keep_mask`. Can you explain?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! These three masks implement the classic BERT masking strategy. Let me explain why we need all three:\n",
"\n",
"**The goal**: When we select a token for masking (with 30% probability), we don't always replace it the same way. Instead:\n",
"\n",
"- **80% of selected tokens** → Replace with `[MASK]` token (mask_mask)\n",
"- **10% of selected tokens** → Replace with random token (random_mask) \n",
"- **10% of selected tokens** → Keep original token (keep_mask)\n",
"\n",
"**Why this strategy?** If we always replaced with `[MASK]`, the model would only learn to predict masked tokens, but during inference there are no `[MASK]` tokens! \n",
"\n",
"By sometimes keeping the original token or using random tokens, the model learns to be robust and can handle any input.\n",
"\n",
"Looking at your example: tokens at positions 4 and 8 were selected for masking, and both got the `mask_mask` treatment (replaced with 0).\n",
"\n",
"Does this make sense? Would you like to see what happens if we adjust the probabilities?"
]
},
{
"cell_type": "markdown",
"id": "25aa892b",
"metadata": {
"input_tokens": 25,
"output_tokens": 171,
"solveit_ai": true
},
"source": [
"`mlm_masking` is applied to packed batches of sequences, right? \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Yes, exactly! The `mlm_masking` function is applied to the packed sequences after they've been assembled into batches.\n",
"\n",
"The flow is:\n",
"1. Individual sequences get packed together into fixed-length packed sequences\n",
"2. Multiple packed sequences form a batch \n",
"3. Then `mlm_masking` is applied to the entire batch at once\n",
"\n",
"So when you see `batch` being passed to `mlm_masking`, that batch contains multiple packed sequences, each potentially containing several original sequences concatenated together.\n",
"\n",
"Does this help clarify when the masking happens in the pipeline?"
]
},
{
"cell_type": "markdown",
"id": "0041f65c",
"metadata": {
"heading_collapsed": true,
"input_tokens": 6
},
"source": [
"### find_best_fit"
]
},
{
"cell_type": "markdown",
"id": "105098d0",
"metadata": {
"hidden": 1,
"input_tokens": 7,
"output_tokens": 240,
"solveit_ai": true
},
"source": [
"Okay let's move on.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great! Now let's look at the next piece - the core packing algorithm. Here's the numba-optimized function that finds the best fit for sequences:\n",
"\n",
"```python\n",
"@njit\n",
"def find_best_fit(remaining_spaces, seq_len):\n",
" valid_spaces = seq_len <= remaining_spaces\n",
" if np.any(valid_spaces):\n",
" valid_space_sizes = remaining_spaces[valid_spaces]\n",
" best_fit_idx = np.argmin(valid_space_sizes)\n",
" return np.arange(len(remaining_spaces))[valid_spaces][best_fit_idx]\n",
" return -1\n",
"```\n",
"\n",
"This function finds the smallest available space that can fit a sequence of `seq_len` tokens. The `@njit` decorator makes it run much faster.\n",
"\n",
"Would you like to try this function with a simple example to see how it works?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "14b1c9d3",
"metadata": {
"hidden": 1,
"input_tokens": 112
},
"outputs": [],
"source": [
"@njit\n",
"def find_best_fit(remaining_spaces, seq_len):\n",
" valid_spaces = seq_len <= remaining_spaces\n",
" if np.any(valid_spaces):\n",
" valid_space_sizes = remaining_spaces[valid_spaces]\n",
" best_fit_idx = np.argmin(valid_space_sizes)\n",
" return np.arange(len(remaining_spaces))[valid_spaces][best_fit_idx]\n",
" return -1"
]
},
{
"cell_type": "markdown",
"id": "c2d97a63",
"metadata": {
"hidden": 1,
"input_tokens": 4,
"output_tokens": 145,
"solveit_ai": true
},
"source": [
"yes plz\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! Let's create a simple example. Imagine you have 3 packed sequences with different amounts of remaining space:\n",
"\n",
"```python\n",
"remaining_spaces = np.array([100, 50, 200]) # space left in each packed sequence\n",
"seq_len = 75 # we want to fit a sequence of length 75\n",
"```\n",
"\n",
"Try running this first, then we can call `find_best_fit` to see which packed sequence it chooses. What do you think will happen?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "97c1a613",
"metadata": {
"hidden": 1,
"input_tokens": 60
},
"outputs": [],
"source": [
"remaining_spaces = np.array([100, 50, 200]) # space left in each packed sequence\n",
"seq_len = 75 # we want to fit a sequence of length 75"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6f30d4a8",
"metadata": {
"hidden": 1,
"input_tokens": 15,
"output_tokens": 19
},
"outputs": [
{
"data": {
"text/plain": [
"0"
]
},
"execution_count": 27,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"find_best_fit(remaining_spaces, seq_len)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8215f08a",
"metadata": {
"hidden": 1,
"input_tokens": 16,
"output_tokens": 31
},
"outputs": [
{
"data": {
"text/plain": [
"array([ True, False, True])"
]
},
"execution_count": 29,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"valid_spaces = seq_len <= remaining_spaces\r\n",
"valid_spaces"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7d48213f",
"metadata": {
"hidden": 1,
"input_tokens": 19,
"output_tokens": 28
},
"outputs": [
{
"data": {
"text/plain": [
"array([100, 200])"
]
},
"execution_count": 31,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"valid_space_sizes = remaining_spaces[valid_spaces]\r\n",
"valid_space_sizes"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "908d8229",
"metadata": {
"hidden": 1,
"input_tokens": 21,
"output_tokens": 25
},
"outputs": [
{
"data": {
"text/plain": [
"np.int64(0)"
]
},
"execution_count": 33,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"best_fit_idx = np.argmin(valid_space_sizes)\r\n",
"best_fit_idx"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6094db08",
"metadata": {
"hidden": 1,
"input_tokens": 21,
"output_tokens": 25
},
"outputs": [
{
"data": {
"text/plain": [
"np.int64(0)"
]
},
"execution_count": 35,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"np.arange(len(remaining_spaces))[valid_spaces][best_fit_idx]"
]
},
{
"cell_type": "markdown",
"id": "2ec5b492",
"metadata": {
"hidden": 1,
"input_tokens": 22,
"output_tokens": 1053,
"solveit_ai": true
},
"source": [
"Okay I understand how it works. How does `njit` work?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! `@njit` is a Numba decorator that compiles Python functions to fast machine code. Here's how it works:\n",
"\n",
"When a call is made to a Numba decorated function it is compiled to machine code \"just-in-time\" for execution and all or part of your code can subsequently run at native machine code speed! [^1] [^2]\n",
"\n",
"The key points:\n",
"\n",
"1. **Just-in-Time Compilation**: Numba reads the Python bytecode for a decorated function and combines this with information about the types of the input arguments to the function [^3] [^4]\n",
"\n",
"2. **`@njit` vs `@jit`**: numba.njit is a function decorator that is similar to numba.jit, but it has stricter requirements for the types of input and output that the function can accept and return. This means that numba.njit may be more restrictive in the types of functions it can compile, but it may also be faster and more memory-efficient than numba.jit. [^5]\n",
"\n",
"3. **Nopython Mode**: The behaviour of the nopython compilation mode is to essentially compile the decorated function so that it will run entirely without the involvement of the Python interpreter [^6] [^7]\n",
"\n",
"So `@njit` makes the `find_best_fit` function run much faster by converting it to optimized machine code instead of interpreted Python!\n",
"\n",
"Would you like to see how much faster it is, or shall we move on to the next piece of code?\n",
"\n",
"[^1]: https://numba.pydata.org/numba-doc/dev/user/5minguide.html\n",
"\t\"When a call is made to a Numba decorated function it is compiled to machine code “just-in-time” for execution and all or part of your code can subsequ...\"\n",
"\n",
"[^2]: https://numba.readthedocs.io/en/stable/user/5minguide.html\n",
"\t\"When a call is made to a Numba-decorated function it is compiled to machine code “just-in-time” for execution and all or part of your code can subsequ...\"\n",
"\n",
"[^3]: https://numba.pydata.org/numba-doc/dev/user/5minguide.html\n",
"\t\"Numba reads the Python bytecode for a decorated function and combines this with information about the types of the input arguments to the function. \"\n",
"\n",
"[^4]: https://numba.readthedocs.io/en/stable/user/5minguide.html\n",
"\t\"Numba reads the Python bytecode for a decorated function and combines this with information about the types of the input arguments to the function. \"\n",
"\n",
"[^5]: https://safjan.com/difference-between-numba-jit-and-njit-decorators/\n",
"\t\"numba.njit is a function decorator that is similar to numba.jit, but it has stricter requirements for the types of input and output that the function ...\"\n",
"\n",
"[^6]: https://numba.pydata.org/numba-doc/dev/user/5minguide.html\n",
"\t\"The behaviour of the nopython compilation mode is to essentially compile the decorated function so that it will run entirely without the involvement o...\"\n",
"\n",
"[^7]: https://numba.readthedocs.io/en/stable/user/5minguide.html\n",
"\t\"The behaviour of the nopython compilation mode is to essentially compile the decorated function so that it will run entirely without the involvement o...\""
]
},
{
"cell_type": "markdown",
"id": "a1df85cf",
"metadata": {
"hidden": 1,
"input_tokens": 27,
"output_tokens": 213,
"solveit_ai": true
},
"source": [
"Can you give me an example to try out how fast it is? I am curious.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great idea! Let's create a simple speed comparison. Here's a test you can try:\n",
"\n",
"```python\n",
"import time\n",
"\n",
"# Create larger test data\n",
"big_remaining_spaces = np.random.randint(50, 500, size=10000)\n",
"test_seq_len = 100\n",
"\n",
"# Time the njit version\n",
"start = time.time()\n",
"for _ in range(1000):\n",
" find_best_fit(big_remaining_spaces, test_seq_len)\n",
"njit_time = time.time() - start\n",
"\n",
"print(f\"Numba @njit time: {njit_time:.4f} seconds\")\n",
"```\n",
"\n",
"Try this first, then we can create a regular Python version to compare. What do you think the timing will be?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "40ff4c0f",
"metadata": {
"hidden": 1,
"input_tokens": 144,
"output_tokens": 39
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Numba @njit time: 0.0436 seconds\n"
]
}
],
"source": [
"import time\n",
"\n",
"# Create larger test data\n",
"big_remaining_spaces = np.random.randint(50, 500, size=10000)\n",
"test_seq_len = 100\n",
"\n",
"# Time the njit version\n",
"start = time.time()\n",
"for _ in range(1000):\n",
" find_best_fit(big_remaining_spaces, test_seq_len)\n",
"njit_time = time.time() - start\n",
"\n",
"print(f\"Numba @njit time: {njit_time:.4f} seconds\")"
]
},
{
"cell_type": "markdown",
"id": "ccb4d387",
"metadata": {
"hidden": 1,
"input_tokens": 7,
"output_tokens": 183,
"solveit_ai": true
},
"source": [
"Let's try the regular one\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! Now let's create the regular Python version without numba:\n",
"\n",
"```python\n",
"def find_best_fit_regular(remaining_spaces, seq_len):\n",
" valid_spaces = seq_len <= remaining_spaces\n",
" if np.any(valid_spaces):\n",
" valid_space_sizes = remaining_spaces[valid_spaces]\n",
" best_fit_idx = np.argmin(valid_space_sizes)\n",
" return np.arange(len(remaining_spaces))[valid_spaces][best_fit_idx]\n",
" return -1\n",
"```\n",
"\n",
"Try this, then we can time it the same way. What do you predict - will it be faster or slower than 0.0893 seconds?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1340ad2e",
"metadata": {
"hidden": 1,
"input_tokens": 108
},
"outputs": [],
"source": [
"def find_best_fit_regular(remaining_spaces, seq_len):\n",
" valid_spaces = seq_len <= remaining_spaces\n",
" if np.any(valid_spaces):\n",
" valid_space_sizes = remaining_spaces[valid_spaces]\n",
" best_fit_idx = np.argmin(valid_space_sizes)\n",
" return np.arange(len(remaining_spaces))[valid_spaces][best_fit_idx]\n",
" return -1"
]
},
{
"cell_type": "markdown",
"id": "9b05a04f",
"metadata": {
"hidden": 1,
"input_tokens": 6,
"output_tokens": 157,
"solveit_ai": true
},
"source": [
"Okay let's try it\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great! Now let's time the regular version:\n",
"\n",
"```python\n",
"# Time the regular Python version\n",
"start = time.time()\n",
"for _ in range(1000):\n",
" find_best_fit_regular(big_remaining_spaces, test_seq_len)\n",
"regular_time = time.time() - start\n",
"\n",
"print(f\"Regular Python time: {regular_time:.4f} seconds\")\n",
"print(f\"Speedup: {regular_time / njit_time:.1f}x faster with numba\")\n",
"```\n",
"\n",
"What do you think the speedup will be?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f338b91b",
"metadata": {
"hidden": 1,
"input_tokens": 121,
"output_tokens": 54
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Regular Python time: 0.1084 seconds\n",
"Speedup: 2.5x faster with numba\n"
]
}
],
"source": [
"# Time the regular Python version\n",
"start = time.time()\n",
"for _ in range(1000):\n",
" find_best_fit_regular(big_remaining_spaces, test_seq_len)\n",
"regular_time = time.time() - start\n",
"\n",
"print(f\"Regular Python time: {regular_time:.4f} seconds\")\n",
"print(f\"Speedup: {regular_time / njit_time:.1f}x faster with numba\")"
]
},
{
"cell_type": "markdown",
"id": "716e28c2",
"metadata": {
"hidden": 1,
"input_tokens": 4,
"output_tokens": 106,
"solveit_ai": true
},
"source": [
"Very nice.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Nice! A 1.6x speedup is pretty good for such a simple change. The speedup can be even more dramatic with more complex numerical operations.\n",
"\n",
"Ready to look at the next piece of the sequence packer? We could explore the scheduler classes next, or jump to the main packing logic in `_create_batch`. What interests you more?"
]
},
{
"cell_type": "markdown",
"id": "24c1ce73",
"metadata": {
"heading_collapsed": true,
"input_tokens": 4
},
"source": [
"### creat_batch"
]
},
{
"cell_type": "markdown",
"id": "6ffc7cdb",
"metadata": {
"hidden": 1,
"input_tokens": 10,
"output_tokens": 274,
"solveit_ai": true
},
"source": [
"Let's jump into the main logic!\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! The main packing logic is in the `_create_batch` method of `GreedyBestFitSequencePacker`. It's quite dense, so let's break it into smaller pieces.\n",
"\n",
"Here's the setup part:\n",
"\n",
"```python\n",
"def _create_batch(self):\n",
" if self.batch_size_scheduler:\n",
" self.out_batch_size = self.batch_size_scheduler(self._token_count)\n",
"\n",
" batch = np.full((self.out_batch_size, self.out_pseq_len), self.pad_token_id, dtype=np.int64)\n",
" seq_counts = np.zeros(self.out_batch_size, dtype=np.int32)\n",
" cum_seq_lens = [[0] for _ in range(self.out_batch_size)]\n",
" remaining_spaces = np.full((self.out_batch_size,), self.out_pseq_len, dtype=np.int32)\n",
" temp_buffer = []\n",
"```\n",
"\n",
"This initializes the arrays for packing. Can you guess what each array is for based on their names?"
]
},
{
"cell_type": "markdown",
"id": "ecb28f9f",
"metadata": {
"hidden": 1,
"input_tokens": 151,
"output_tokens": 198,
"solveit_ai": true
},
"source": [
"Here is my guess:\r\n",
"\r\n",
"- `batch`: creates an outut batch with size `(self.out_batch_size, self.out_pseq_len)`.\r\n",
"- `seq_counts`: number of sequences in each batch with size `self.out_batch_size`.\r\n",
"- `cum_seq_lens`: cumulative sequences lengths, which includes information about when each sequence ends starting from 0.\r\n",
"- `remaining_spaces`: Contains information about how much empty space is in each batch.\r\n",
"- `temp_buffer`: Temporary buffer? For packing?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Excellent guesses! You're spot-on with most of them:\n",
"\n",
"- ✅ `batch`: Correct - the output tensor filled with pad tokens initially\n",
"- ✅ `seq_counts`: Correct - tracks how many sequences are packed in each row\n",
"- ✅ `cum_seq_lens`: Correct - tracks where each sequence starts/ends within each packed sequence\n",
"- ✅ `remaining_spaces`: Correct - how much space is left in each packed sequence\n",
"- ✅ `temp_buffer`: Good guess! It's for sequences that don't fit in the current batch and need to be saved for the next one\n",
"\n",
"You really understand the structure! Ready to see the main packing loop where the magic happens?"
]
},
{
"cell_type": "markdown",
"id": "80012721",
"metadata": {
"hidden": 1,
"input_tokens": 3,
"output_tokens": 385,
"solveit_ai": true
},
"source": [
"yea\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great! Here's the main packing loop:\n",
"\n",
"```python\n",
"while True:\n",
" # Check if buffer has more items, and if not replenish\n",
" if not self.buffer:\n",
" items_to_fetch = self.buffer_size - len(temp_buffer)\n",
" items_added = self._fill_buffer(items_to_fetch)\n",
" if items_added == 0:\n",
" break\n",
"\n",
" seq = self.buffer.popleft()\n",
" seq_len = len(seq)\n",
"\n",
" # Find the best fit (smallest space that can accommodate the sequence)\n",
" best_fit_idx = find_best_fit(remaining_spaces, seq_len)\n",
" if best_fit_idx != -1:\n",
" # Pack the sequence\n",
" end_pos = self.out_pseq_len - remaining_spaces[best_fit_idx]\n",
" batch[best_fit_idx, end_pos : end_pos + seq_len] = seq\n",
" seq_counts[best_fit_idx] += 1\n",
" remaining_spaces[best_fit_idx] -= seq_len\n",
" cum_seq_lens[best_fit_idx].append(cum_seq_lens[best_fit_idx][-1] + seq_len)\n",
" else:\n",
" # Can't fit the sequence, save for next batch\n",
" temp_buffer.append(seq)\n",
"```\n",
"\n",
"Can you trace through what happens when a sequence fits vs when it doesn't fit?"
]
},
{
"cell_type": "markdown",
"id": "dc16b872",
"metadata": {
"hidden": 1,
"input_tokens": 12,
"output_tokens": 186,
"solveit_ai": true
},
"source": [
"What does `self.buffer` do?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great question! `self.buffer` is the internal storage that holds individual sequences waiting to be packed.\n",
"\n",
"It's a `deque` (double-ended queue) that acts like a line at a store:\n",
"- New sequences get added to the back when `_fill_buffer()` pulls from the source data\n",
"- The packer takes sequences from the front with `self.buffer.popleft()`\n",
"- If a sequence doesn't fit in the current batch, it goes to `temp_buffer` and gets put back at the front later\n",
"\n",
"Think of it as the \"waiting room\" for sequences before they get packed together. Does this make sense?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8faa2e9",
"metadata": {
"hidden": 1,
"input_tokens": 7,
"output_tokens": 21
},
"outputs": [
{
"data": {
"text/plain": [
"deque([])"
]
},
"execution_count": 43,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"buffer = deque()\r\n",
"buffer"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f641a399",
"metadata": {
"hidden": 1,
"input_tokens": 3,
"output_tokens": 19
},
"outputs": [
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 47,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"not buffer"
]
},
{
"cell_type": "markdown",
"id": "9e34262f",
"metadata": {
"hidden": 1,
"input_tokens": 19,
"output_tokens": 106,
"solveit_ai": true
},
"source": [
"It fills buffer only when the `self.buffer` is empty?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly! You've got it. The buffer only gets refilled when it's completely empty (`if not self.buffer:`).\n",
"\n",
"This is a simple but effective strategy - it processes all the sequences currently in the buffer before fetching more from the source data.\n",
"\n",
"Would you like to see what happens in the `_fill_buffer()` method, or continue tracing through the packing logic?"
]
},
{
"cell_type": "markdown",
"id": "ff227387",
"metadata": {
"hidden": 1,
"input_tokens": 193,
"output_tokens": 193,
"solveit_ai": true
},
"source": [
"Let me focus on this then I will move on to how it fills buffer later. \r\n",
"\r\n",
"We pop the first sequence from the `self.buffer` and check whether we can fit in the batches. If we can, `best_fit_idx` is not -1 and we pack the sequence. We will put it inside of remaining_spaces[`best_fit_idx`]. So we get `end_pos` information, put it in the `batch`, and update `seq_counts`, `remaining_spaces[best_fit_idx]`, `cum_seq_lens[best_fit_idx]`. If the sequence has no room, put it in the `temp_buffer`.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! You've understood the core packing logic exactly right. \n",
"\n",
"One small clarification: we don't put the sequence \"inside\" `remaining_spaces[best_fit_idx]` - that array just tracks how much space is left. The sequence actually goes into `batch[best_fit_idx, end_pos : end_pos + seq_len]`.\n",
"\n",
"But yes, you've got the flow:\n",
"1. Pop sequence from buffer\n",
"2. Find best fit → get index or -1\n",
"3. If fits: pack it and update all the tracking arrays\n",
"4. If doesn't fit: save for next batch\n",
"\n",
"Ready to see what happens after the main loop finishes?"
]
},
{
"cell_type": "markdown",
"id": "c40f2ec8",
"metadata": {
"hidden": 1,
"input_tokens": 10,
"output_tokens": 234,
"solveit_ai": true
},
"source": [
"yep. Let's move on.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! Here's what happens after the main loop - the cleanup and return:\n",
"\n",
"```python\n",
"# Add any sequences we skipped back to the start of the buffer\n",
"self.buffer.extendleft(temp_buffer)\n",
"\n",
"if np.all(seq_counts > 0):\n",
" self._seqs_emitted += np.sum(seq_counts)\n",
" for x in cum_seq_lens:\n",
" if x[-1] != self.out_pseq_len:\n",
" x.append(self.out_pseq_len)\n",
" return batch, cum_seq_lens\n",
"else:\n",
" # If we can't form a full batch, we return None to signal the end\n",
" return None\n",
"```\n",
"\n",
"Can you figure out what the `if np.all(seq_counts > 0):` check is doing and why it might return `None`?"
]
},
{
"cell_type": "markdown",
"id": "41d6e10b",
"metadata": {
"hidden": 1,
"input_tokens": 48,
"output_tokens": 175,
"solveit_ai": true
},
"source": [
"I think it checks whether we created all the batches by packing sequences. If there was not enough sequences left to create all the batches, it returns `None`.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly right! \n",
"\n",
"The `np.all(seq_counts > 0)` checks if every packed sequence has at least one sequence in it. If some packed sequences are still empty (seq_count = 0), it means we ran out of input data and couldn't fill a complete batch.\n",
"\n",
"Returning `None` signals to the iterator that we're done - no more batches can be created. This is how the sequence packer knows when to stop generating batches.\n",
"\n",
"You've really grasped the core logic! Ready to look at `_fill_buffer()` now, or would you like to explore a different part?"
]
},
{
"cell_type": "markdown",
"id": "228f4660",
"metadata": {
"hidden": 1,
"input_tokens": 151,
"output_tokens": 205,
"solveit_ai": true
},
"source": [
"I want to finish this part. We save sequences that didn't fit into the buffer to the left so they are popped first in the next batch. And if it has all the packed sequences to return, it adds up all the counts of sequences into `self._seqs_emitted`. Before it returns the batches with `cum_seq_lens`, `cum_seq_lens` needs to add `self.out_pseq_len` at the end. I think the last part of the batch contains padding tokens.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Excellent analysis! You've understood it perfectly:\n",
"\n",
"1. ✅ Sequences that didn't fit go to the **left** of the buffer (`extendleft`) so they're processed first next time\n",
"2. ✅ `self._seqs_emitted` tracks the total sequences that have been successfully packed\n",
"3. ✅ Adding `self.out_pseq_len` at the end of `cum_seq_lens` marks where the packed sequence ends (including any padding)\n",
"\n",
"You're absolutely right about the padding - if the last sequence doesn't fill the entire `out_pseq_len`, the remaining space gets filled with pad tokens.\n",
"\n",
"Great job working through this complex logic! Ready to look at `_fill_buffer()` now?"
]
},
{
"cell_type": "markdown",
"id": "575a19e5",
"metadata": {
"heading_collapsed": true,
"input_tokens": 4
},
"source": [
"### fill_buffer"
]
},
{
"cell_type": "markdown",
"id": "4e86b85c",
"metadata": {
"hidden": 1,
"input_tokens": 1,
"output_tokens": 285,
"solveit_ai": true
},
"source": [
"yeah\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great! Here's the `_fill_buffer()` method:\n",
"\n",
"```python\n",
"def _fill_buffer(self, max_items_to_add=float(\"inf\")) -> int:\n",
" items_added = 0\n",
" while (self.buffer_size - len(self.buffer)) > self.src_batch_size:\n",
" try:\n",
" if max_items_to_add < float(\"inf\"):\n",
" if (items_added + self.src_batch_size) > max_items_to_add:\n",
" break\n",
" incoming_batch = next(self.src_iterator)\n",
" assert len(incoming_batch) <= self.src_batch_size\n",
" for item in incoming_batch:\n",
" if len(item[\"input_ids\"]) > 0: # ignore empty sequences\n",
" self.buffer.append(item[\"input_ids\"])\n",
" items_added += 1\n",
" self._seqs_consumed += 1\n",
" except StopIteration:\n",
" break\n",
" return items_added\n",
"```\n",
"\n",
"What do you think the `while` condition is checking for?"
]
},
{
"cell_type": "markdown",
"id": "747f7443",
"metadata": {
"hidden": 1,
"input_tokens": 91,
"output_tokens": 124,
"solveit_ai": true
},
"source": [
"`self.src_batch_size` is incoming batch size, `self.buffer_size` is the maximum size of buffer for seqs, and `len(self.buffer)` is how many seqs in the buffer currently. So, the while loop is checking for the room in the buffer for the incoming batch. \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly right! The condition `(self.buffer_size - len(self.buffer)) > self.src_batch_size` ensures there's enough room in the buffer for at least one full incoming batch.\n",
"\n",
"It's being conservative - it won't pull a new batch unless there's definitely space for all of it. This prevents the buffer from overflowing.\n",
"\n",
"Can you trace through what happens inside the loop when it successfully gets a batch from `self.src_iterator`?"
]
},
{
"cell_type": "markdown",
"id": "8274052f",
"metadata": {
"hidden": 1,
"input_tokens": 94,
"output_tokens": 220,
"solveit_ai": true
},
"source": [
"After grabbing a batch, it checks whether the size of the batch is within the `src_batch_size`. If the check passes, we add the `item[\"input_ids\"]` into the buffer, increment `items_added` and `self._seqs_consumed`. It also skips length 0 sequence. \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! You've got it exactly right:\n",
"\n",
"1. ✅ Asserts the batch size is valid (`<= self.src_batch_size`)\n",
"2. ✅ Extracts `item[\"input_ids\"]` (the actual token sequences) and adds to buffer\n",
"3. ✅ Increments counters: `items_added` (for this call) and `self._seqs_consumed` (total)\n",
"4. ✅ Skips empty sequences (`len(item[\"input_ids\"]) > 0`)\n",
"\n",
"You've really mastered the core sequence packer logic! The data flows from `src_iterator` → `buffer` → packed batches.\n",
"\n",
"What would you like to explore next? The scheduler classes, or try creating a simple example to see the whole thing in action?"
]
},
{
"cell_type": "markdown",
"id": "96ec9ec8",
"metadata": {
"input_tokens": 6
},
"source": [
"## Just random thoughts"
]
},
{
"cell_type": "markdown",
"id": "7cfa9c3e",
"metadata": {
"input_tokens": 106,
"output_tokens": 234,
"solveit_ai": true
},
"source": [
"Why does it track `self._seqs_consumed`? The comment says it is to report packing efficiency. My guess is we can compare how many batches of sequences are in the original training set to `seqs_consumed` and `seqs_emitted`. If there is low `seqs_emitted`, it is not packing efficiently. \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Excellent insight! You're absolutely right about the purpose.\n",
"\n",
"`seqs_consumed` vs `seqs_emitted` tells you about packing efficiency:\n",
"\n",
"- `seqs_consumed`: Total sequences read from the source\n",
"- `seqs_emitted`: Total sequences that made it into output batches\n",
"\n",
"If they're very different, it means sequences are getting \"stuck\" in the buffer and not being packed efficiently. \n",
"\n",
"For example, if you have many very long sequences that can't fit anywhere, `seqs_consumed` keeps growing but `seqs_emitted` stays low.\n",
"\n",
"You could calculate packing efficiency as `seqs_emitted / seqs_consumed` - closer to 1.0 means better packing!\n",
"\n",
"Does this match your understanding of how you'd use these metrics?"
]
},
{
"cell_type": "markdown",
"id": "9217d379",
"metadata": {
"input_tokens": 43,
"output_tokens": 223,
"solveit_ai": true
},
"source": [
"What can I do if the packing is not efficient? I think I can change the rng value or increase the `out_pseq_len`. \n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great thinking! You've identified some options, but let me help you understand which would actually help:\n",
"\n",
"**Your suggestions:**\n",
"- **RNG value**: Won't help much - the randomness doesn't affect how sequences fit together\n",
"- **Increase `out_pseq_len`**: ✅ Good idea! Longer packed sequences give more room to fit variable-length sequences\n",
"\n",
"**Other strategies that could help:**\n",
"- **Increase `buffer_size`**: More sequences to choose from = better chance of finding good fits\n",
"- **Sort sequences by length**: Group similar lengths together for better packing\n",
"- **Adjust the source data**: Filter out extremely long sequences that rarely fit\n",
"\n",
"Which of these makes the most sense to you for improving packing efficiency?"
]
},
{
"cell_type": "markdown",
"id": "c181f85b",
"metadata": {
"input_tokens": 13,
"output_tokens": 97,
"solveit_ai": true
},
"source": [
"oh yeah RNG value just changes masking behavior.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly! The RNG only affects which tokens get masked, not how sequences get packed together.\n",
"\n",
"The packing efficiency really comes down to how well different sequence lengths can fit together, like a puzzle. \n",
"\n",
"Which of the other strategies I mentioned seems most practical to you - increasing buffer size, sorting by length, or filtering extreme lengths?"
]
},
{
"cell_type": "markdown",
"id": "26ae3bb2",
"metadata": {
"input_tokens": 60,
"output_tokens": 198,
"solveit_ai": true
},
"source": [
"To use sequence packer with a causal model, I just have to switch from `mlm_masking` to another function that uses a roll on sequences to create labels and return them, right?\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Exactly right! For causal models, you'd replace `mlm_masking` with something that shifts the sequences to create the \"next token prediction\" labels.\n",
"\n",
"The typical approach is:\n",
"- Input: `[token1, token2, token3, token4]`\n",
"- Labels: `[token2, token3, token4, <end_token>]`\n",
"\n",
"So labels are just the input shifted by one position. Much simpler than the BERT-style masking!\n",
"\n",
"The rest of the sequence packer (the packing logic, cu_seqlens, etc.) would work exactly the same way.\n",
"\n",
"Have you thought about how you'd handle the label creation for the packed sequences?"
]
},
{
"cell_type": "markdown",
"id": "627d2e97",
"metadata": {
"input_tokens": 18,
"output_tokens": 151,
"solveit_ai": true
},
"source": [
"yea I can roll once on the seqeunce.\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Perfect! Yes, a simple roll/shift operation would work great. \n",
"\n",
"One thing to consider: with packed sequences, you'll need to be careful about the boundaries between different sequences. You wouldn't want the label for the last token of sequence A to be the first token of sequence B.\n",
"\n",
"The `cu_seqlens` information would help you handle those boundaries properly when creating the shifted labels.\n",
"\n",
"Have you worked with causal language modeling before, or would this be your first time adapting the packer for it?"
]
},
{
"cell_type": "markdown",
"id": "60fbd7ce",
"metadata": {
"input_tokens": 7,
"output_tokens": 108,
"solveit_ai": true
},
"source": [
"I am familiar with it\n",
"\n",
"##### 🤖Reply🤖<!-- SOLVEIT_SEPARATOR_7f3a9b2c -->\n",
"\n",
"Great! Then you already know the challenges with sequence boundaries and attention masking.\n",
"\n",
"You've done a fantastic job understanding how the ModernBERT sequence packer works - from the core packing algorithm to the masking strategy to the efficiency tracking.\n",
"\n",
"Is there anything else about the sequence packer you'd like to explore, or are you ready to start experimenting with it on your own?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "85220612",
"metadata": {
"input_tokens": 4
},
"outputs": [],
"source": [
"torch.roll?"
]
}
],
"metadata": {
"solveit_dialog_mode": "learning",
"solveit_ver": 2
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment