Ian’s Blog

Category: posts

  • Byte-Latent Transformers

    Who needs a tokenizer anyway!

    [2412.09871] Byte Latent Transformer: Patches Scale Better Than Tokens

    This paper, from back in December last year, presents an interesting approach to handling raw byte sequences in LLMs without relying on tokenization.

    Vocab sizes for tokenizers have gone up over the last couple of years with attendant gains in usefulness, but this remains a particularly hand-tuned number in the training process. BLT proposes a method that processes raw UTF-8 byte sequences directly, leveraging a dynamic patching mechanism to group bytes into variable-length patches based on entropy.

    Higher-entropy regions receive more attention and shorter patches, while lower-entropy regions can be processed more efficiently.

    There are conceptually three levels of processing:

    • Local Encoder: A small transformer stack encodes raw byte sequences into higher level representations, which are then structured into patches.
    • Latent Global Transformer: A standard large transformer model operating on patch-level representations
    • Local Decoder: The encoded patches are decoded back into byte sequences, using a cross-attention mechanism to reconstruct text.

    In the paper they show they can achieve parity in pretraining with a traditional tokenized approach in llama for similar parameter count, while being more robust and offering some inference time performance gains. The patching approach allows for allocating compute where needed most.

    Retrofitting existing models

    One of the ideas I found most interesting is starting with a traditionally pretrained model. The paper discusses using the main transformer layers from Llama and training the byte latent approach successfully.

    I gave the approach a go with a simplified local encoder, entropy and patching approach, and took the transformer layers from Qwen 2.5 3B, a strong model that could still be trained locally (no corporate resources were harmed, etc).

    The basic approach was replacing the tokenizer, adding a small transformer and patch pooling based on a local entropy measure to generate patches, then cross-attending in some of the Qwen layers. Its training a new encoder while leveraging Qwen for the backbone of the global transformer and adding new cross-attention params to make it also the decoder, with the embedding layers at each end chopped off – so a significant domain shift. For inference I leverage the same patch generation process to try and generate effective tokens.

    You can find my Torchtune recipe on GitHub, running through the Alpaca dataset. Thus far I’m still training so while loss is improving, I have no idea whether it will turn into something useful. The fact that there is something trainable is fun though, and I have hopes that this kind of technique will lead to some breakthroughs in tokenizer-free models in the future!

  • On career growth

    Particularly at the giant tech firms where there are many, many smart people, folks look at success and try to copy it. They set themselves up for promotions against the definition of what is expected at the next level. But those expectations describe an average, not a person. Real people are spiky.

    No one is great at everything, in every situation. Some of the most successful can create situations where they can use their strengths. They shape the work to fit them. Most of us don’t get that luxury, but we can often pick where to play. I’ve taken on projects and roles that looked reasonable, but I knew weren’t a fit, and the results have been from poor-to-fair, never great. That mismatch is costly.

    Choosing projects, teams, or roles is more significant than choosing how to work on then. They decide whether you work with your strengths or against them. Lean into strengths, lean into things that bring you energy and satisfaction. That doesn’t mean staying comfortable. It means knowing where you do your best work and pushing those capabilities or abilities even further, rather than attempting to contort yourself to a generic level N+1.

  • GRPO & Verifiable Rewards

    GRPO (Group Relative Policy Optimization) is an RL technique originally proposed in the DeepSeekMath paper. Instead of using a full-blown value network like PPO does, GRPO samples a group of completions for a given prompt and then computes a relative (normalized) reward for each output. The rewards are “verifiable” because they come from checking the final answer against ground truth and confirming. E.g. does the response follow the expected format (i.e. a <think>…</think> block for reasoning and an <answer>…</answer> block for the solution) and is the answer accurate against a predetermined fact. Not every problem fits this model, but there are a bunch that do, including math reasoning with the GSM8K dataset of grade-school math word problems. These look like this:

    “Julie is reading a 120-page book. Yesterday, she was able to read 12 pages and today, she read twice as many pages as yesterday. If she wants to read half of the remaining pages tomorrow, how many pages should she read?”

    How Does the Training Work?

    1. Sampling Completions: For each prompt, the model generates a group of candidate completions. These are produced in inference mode (gradients aren’t collected) using a KV cache for speed (or a dedicated inference engine like VLLM)
    2. Verifiable Reward Calculation: Each completion is scored between 0 and 1—rewarding outputs that follow the prescribed format and yield the correct answer.
    3. Forward Pass for Gradients: Both the “policy” (the model being tuned) and a reference (typically the base, unmodified model) are used for a forward pass with the prompt and completions to compute per-token logits and log-probabilities.
    4. Loss and Backwards: The loss is then calculated as a combination of the (group-averaged) reward and a KL divergence term between the tuned model and the baseline, to constrain learning to similar responses. This loss is backpropagated through the policy model based on the earlier forward pass.

    Getting it going in TorchTune

    Over last weekend I hacked up a quick and dirty version of the training loop in the TorchTune, and over a couple of bus rides to Menlo Park cleaned it up into something that could work as a more general recipe(PR). Most of the work goes into the recipe and getting the dataset shaped properly to generate completions. This version—tested on a smaller model (the 1B Llama 3.2 variant, with LoRA)—showed some promising improvements in approach but I didn’t get to the point of having something converge enough to be confident in the overall recipe. In the DeepSeek R1 paper they had discussed trying a smaller model, but found 3B was the lowest they were able to get results on with some of their fine-tuning approaches.

    Luckily for everyone, at around the same time Ariel Kwiatkowski also put together a version that included distributed device support, making it easier to experiment on bigger models. This PR is more modular, and I’m excited to see it refined and landed so the recipe is widely available!

    There’s a growing energy around tools like torchtune, and it’s exciting to see how easy it is to “hack on” these ideas. It’s also great to see the techniques show up in other libraries, like HuggingFace’s TRL, which is being used as part of the OpenR1 replication effort!

  • Gradient Accumulation (was) busted

    This weekend I was reading the Tulu v3 paper (link), which offers a deep dive into building robust post-training setups. This is an very good resource for anyone aiming to build a really robust fine-tuning workflows. It covers critical elements like data set selection, synthetic data generation (with example prompts!), strategies for SFT and preference tuning, and various things they struggled with.

    One struggle was an issue with gradient accumulation they ran into where the loss was worse than without it on. The community at large also hit this, and fixed it, thanks to an excellent blog post by Unsloth (link).

    The bug

    Gradient accumulation is a technique used to simulate larger batch sizes by accumulating gradients over several smaller batches before performing a backward pass. This approach is particularly useful for managing memory constraints during training, so comes up a lot when post-training on a biggish model with more limited hardware.

    The problem arises when dealing with sequences of varying lengths within these mini-batches. In standard practice, the loss is calculated and normalized by the number of non-padded (i.e., valid) tokens in each sequence. However, when accumulating gradients across multiple mini-batches, each with different sequence lengths, the naive summation of gradients can lead to an incorrect total loss calculation.

    The discrepancy occurs because the cross-entropy loss function normalizes by the number of valid tokens, and this normalization factor can vary between mini-batches. When these normalized losses are accumulated without proper adjustment, the final loss does not match what would have been obtained using a single large batch. This results in a higher observed loss during training when using gradient accumulation compared to full batch training.

    Daniel and co at unsloth addressed this issue by developing a methodology that ensures the accumulated gradients are correctly scaled, accounting for the varying sequence lengths across mini-batches. This fix aligns the gradient accumulation process more closely with the theoretical foundations of full batch training, leading to more accurate loss calculations and improved training performance.

    Fixes and Workarounds

    Recent updates in both Hugging Face Transformers (pull request) and TorchTune (pull request) offer fixes. And, at least in Evan’s case, a little bit snark:

    In honor of the day the ML community first discovered the fact that (x1 / n1) + (x2 / n2) != (x1 + x2) / (n1 + n2)

    I really like seeing these small but practical problems pop up, and seeing the community rally around to fix them. I missed this when it happened in October, so glad to look back at it now!

  • nonzero_static in PyTorch

    Natalia Gimelshein added cuda support for nonzero_static to PyTorch the other day — its a feature Jax has had for a while that lets you avoid a bit of data dependent annoyingness.

    I most often see nonzero pop up in logs: its underlies a lot of boolean mask operations and torch.where. A downside of torch.nonzero is that we don’t know the size of the returned tensor. The shape is data dependent, which causes pain for the compiler, and when running on an accelerator requires a device to host sync. This can significantly slow down otherwise fast operations.

    nonzero_static overcomes this by allowing you to supply a shape for the output — if the actual result is smaller it is padded, if larger the result is truncated. The PR linked above enables that for CUDA. You can’t seamlessly use it with bool masks or torch.where, but you can easily replace the call with one to nonzero_static.

    To see the difference, imagine we have a big tensor where we have some sense of the output shape, for example searching for a one-hot index like this:

    import torch

# Set device
device = torch.device("cuda:0")

# Generate a large tensor of 0s with exactly one 1 at a specific index on the GPU
size = 100_000_000
large_tensor = torch.zeros(size, device=device, dtype=torch.long)
large_tensor[size // 2] = 1  # Place a single 1 in the middle of the tensor

# Record GPU events to ensure asynchronicity
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

start_event.record()
ones_indices = torch.nonzero(large_tensor)
end_event.record()
torch.cuda.synchronize()

elapsed_time_ms = start_event.elapsed_time(end_event)
print(f"torch.nonzero() execution time: {elapsed_time_ms:.2f} ms")
print("Indices of 1s:", ones_indices.cpu())

    That gives us (on my laptop)

    torch.nonzero() execution time: 177.58 ms
    Indices of 1s: tensor([[50000000]])

    If we modify it to use nonzero_static: 

    import torch

# Set device
device = torch.device("cuda:0")

# Generate a large tensor of 0s with exactly one 1 at a specific index on the GPU
size = 100_000_000
large_tensor = torch.zeros(size, device=device, dtype=torch.long)
large_tensor[size // 2] = 1  # Place a single 1 in the middle of the tensor

# Record GPU events to ensure asynchronicity
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

start_event.record()
ones_indices = torch.nonzero_static(large_tensor, size=1)
end_event.record()
torch.cuda.synchronize()

elapsed_time_ms = start_event.elapsed_time(end_event)
print(f"torch.nonzero_static() execution time: {elapsed_time_ms:.2f} ms")
print("Indices of 1s:", ones_indices.cpu())  

    torch.nonzero_static() execution time: 78.53 ms
    Indices of 1s: tensor([[50000000]])

    A very nice improvement!

  • Engineering Culture at Meta

    A question I’ve been asked recently is how Meta compares to other places I’ve worked, or what makes it different.

    From my conversations and observations, those who disliked working at Meta often cited chaos, short-term focus, and internal politics, while those who liked it called out  autonomy, speed, and the feeling they could work on important projects. To explain that disparity, I refer to three values or themes that shape the work culture.  

    Individuals are responsible for doing impactful work: “impact” is an important concept at Meta, and having a level-appropriate collection of impactful work at performance review time is important for every engineer. If you find yourself in a situation where your impact feels limited, you are generally responsible for exploring ways to address it, or make a change. 

    This means that ICs (individual contributors) at Meta are willing to cross team boundaries to find important work, and will also gravitate towards highly visible projects. They care about how their work is regarded and how it fits in to the wider organization. Internal mobility is fairly easy, so folks will leave teams if they can’t find the right kind of work. Its also reflected in the growth expectations: when I worked at Google and Lyft they also had expectations that IC3s would become IC4s, and IC4s become IC5s (though Google later removed this part), but the timelines were somewhat soft. At Meta, they are firm, and expectations ramp at defined intervals as you approach the boundaries. 

    Practically, this means ICs should expect to identify and collaborate on projects that align with organizational goals and take initiative to push them forward. Managers and leadership provide support, but success heavily depends on individuals ability and desire to chart a path, adapt as needed, and ensure their contributions are visible. In general, there’s a strong bias towards getting things done, getting things out,  “rough consensus and running code”.

    Dave Anderson has written about how much more helpful he found teams at Meta, vs Amazon where there was a lot of horse trading for collaboration. Part of that is driven, I think, by this responsibility for impact. Having another team’s thanks, or enabling results for them, allows you to claim some credit for their impact with relatively little effort. Conversely, intentionally blocking another team can be seen as gatekeeping, which is frowned upon.

    No gatekeeping: Some version of “Move Fast” has been in Meta’s official values for a long time, and the company still operates at a good clip. Part of that is aided by generally making it easy to go make changes wherever they need to be made. One example of this I use with Google folks is OWNERS files. Google and Meta are both monorepo based, but at Google you have sets of services with clear owners, and touching code in another team’s service requires their full blessing. Meta also operates a monorepo, but there is much more fluidity — folks can land changes anywhere they need to. In part this is because the original Meta product, Facebook itself, is a monolith, but there is a deeper cultural aspect here.

    For example, even very senior engineers will very rarely say “no” to something. Instead of outright rejection, feedback is often framed as suggestions or concerns to consider. This encourages risk-taking and innovation, but it also places a significant responsibility on engineers to weigh feedback carefully, address risks , and seek out champions from stakeholders. This is one source of the disconnect between folks who see Meta as a place where it’s ok to take risks and folks who don’t: if you take a risk, it fails, and at PSC (performance review) time someone affirms they called out the problems that occurred and you didn’t take appropriate measures, you will be dinged. I have seen people interpret this kind of feedback as nits or suggestions, rather than weighing it heavily and convincing others that the risk is well managed ahead of time.

    In general, folks are expected to be helpful, to provide guidance to others, and not to put up walls, so it can be a tricky balance when outside teams or other engineers come in and ride roughshod over a team’s plans or projects. As an engineer, escalating misalignments on goals/priorities to management is usually well supported, and as a manager, putting engineers together to get to technical solutions across teams is expected. Enacting hard blocks where one team can’t achieve their goals because another was in the way is less so.

    The heavy dependence on individuals and relationships, particularly for cross-team projects, is another key theme:

    It’s a social company: Somewhat unsurprisingly for a company that started around a social network, Meta is a pretty social company. The internal social network is a firehose of information, and there are deep networks of connections across the company between senior engineers, managers and executives.

    Its important for engineers to talk about their work in order to find folks interested in it, build connections and relationships with them, and have a good sense not just of their org but the universe of organizations that they operate within. A fairly common failure pattern is to build a good relationship with one side of an org and ignore another, developing a sizable blind spot that later comes back to be a problem.

    The official org chart is the secondary and lagging structure at the company. The more important structure is the informal network of relationships that where many things get done, and decisions get made.

    This dynamic can sometimes feel political, which is why some describe Meta this way. While there are large-company politics (it is a large, influential company), for most people, it’s less about traditional power politics and more about navigating cliques and informal networks of folks who have worked together on multiple projects and have mutual trust and respect. The company leans a lot on strong, senior engineers to drive projects to success, and those folks may not report into the org that is officially doing the work, or may be at a lower or higher level in the org chart than you might expect. They will often work by going directly to the people they know to unblock issues, drive important changes, or get alignment on a controversial decisions.

    For big enough changes, org structural changes do follow, but they usually lag rather than lead the work itself.

    What are the downsides and upsides?

    As folks who have had a bad time can attest, Meta can be chaotic. There can be parallel implementations, people can swarm on important projects to the detriment of those trying to work on them, and less impactful projects can end up unowned and passed around. It can feel like information overload, with more being published than you can possibly follow. At the same time, there’s can be an information drought when truly important conversations happen in small, exclusive groups. For example, I’ve seen feedback about a project be shared openly between a small collection of engineers and leaders, without ever clearly reaching the team responsible for developing it.

    The flip side is that this combination of values allows Meta to pivot surprisingly fast. Changes that would have required months at some companies can be kicked off in a day, particularly by very senior, well-connected leaders. Senior ICs can take problem descriptions, and quickly form an idea of who might have thoughts on it, and pull them in. Soon you have a loose group (often later structured into a “v-team” or virtual team) that can quickly align and drive change. The lack of gatekeeping, both technically and culturally, reduces the corporate immune reaction to large changes. The incentive to individual impact encourages folks to jump onboard important things without having to work out if that means a team change, or what their long-term situation might be.

  • A very lazy mental model for parallelisms

    There is a great deal of technical complexity in distributed data parallel, model parallel, tensor parallel, fully-sharded-data-parallel, ZeRO 1/2/3, context parallel etc. To fit ideas into a picture I tend to segment in-practice parallelism into three buckets.

    1. Data parallel
    2. Pipeline parallelism
    3. Everything else

    I will tackle these out of order:

    Data parallelism

    Data parallel is a very convenient training concept because you take multiple copies of the model and give different data to them, do a forward-backward pass, then all-reduce gradients between them i.e., aggregate gradients across all devices so they operate as if they had trained on the unified set of data. This lets you train more quickly by using more devices!

    This all-reduce operation can usually be effectively overlapped with other computations, minimizing the overhead of data parallelism, and the copies of the model can be largely anything: a model running on a single GPU, or a cluster of devices running many parallelisms of their own.

    Because you get to run a whole batch of the computation through the model before needing comms, its relatively network efficient as well, which makes it plausible to do distributed data parallel over more normal network links (than the ones I’m about to mention). 

    Everything else

    Almost everything about these parallelism methods are annoying: either the communication collectives are difficult to overlap with other operations, or they must occur more frequently, leading to increased sensitivity to latency.

    This means you want faster links, which means you need very tight networking like, canonically, NVLink, which is Nvidia’s high-speed interconnect between GPUs. This tends to mean you have tensor parallelism, context parallelism etc. implemented within these fast domains, which traditionally is one host (8 GPUs), but with Blackwell can be quite a lot more (up to 36 or 72). There are a lot of options for efficiently packing the compute and scheduling work between devices in there based on the characteristics of the model, and they form this big set of potential parallelisms.

    Pipeline parallelism 

    When your model unit is too big for the available “fast” networking domain, you try and divide the model into multiple sequential stages (like layers) and process them in parallel, overlapping the computation of one stage with the communication of the next.

    This is painful for all the reasons pipelining in anything is painful but is additionally painful in that you have to write your training code with, effectively, a bunch of if statements for whichever stage it happens to be in. 

    There is some leakage between these buckets (e.g. the first and last stages of the pipeline tend to vary a lot due to dealing with embeddings for input and output), but it’s a reasonable first approximation to treat them as discrete.

    (Or FSDP everything) 

    This does presume quite a lot of scale, as the cost of getting everything working well together is non-trivial. FSDP (preferably FSDP2 in PyTorch at least) directly mixes up between the DP and Everything Else buckets, and works pretty well for most folks up to a decently large (100s+) number of GPUs. So roughly: 

    • Smallish model, lots of data: Distributed data parallel.
    • Medium sized model, medium sized number of GPUs: FSDP2
    • Anything expensive: units of <parallelisms>, chunked into pipelines if the model is too big, copied into multiple DP copies to speed up training.  

  • PyTorch while_loop

    I’ve been following the development of the higher order ops in PyTorch nightlies for a little bit, and got a chance to try out while_loop. The best examples right now are in the tests, but as another, here’s a mandlebrot example: 

    import torch
from torch._higher_order_ops.while_loop import while_loop
import matplotlib.pyplot as plt

def mandelbrot_step(z, c):
    """Performs one iteration of the Mandelbrot sequence."""
    return z**2 + c

def mandelbrot(c, max_iter, threshold):
    """Compute Mandelbrot set membership for a grid of complex numbers."""
    def cond_fn(z, iter_count, mask):
        return torch.any(mask & (iter_count < max_iter))

    def body_fn(z, iter_count, mask):
        z_next = mandelbrot_step(z, c)
        diverged = torch.abs(z_next) > threshold
        mask_next = mask & ~diverged
        iter_count_next = iter_count + mask_next
        return z_next, iter_count_next, mask_next

    # Initialize variables
    z0 = torch.zeros_like(c)
    iter_count = torch.zeros(c.shape, dtype=torch.int32)
    mask = torch.ones(c.shape, dtype=torch.bool)  # All points start as candidates
    final_state = while_loop(cond_fn, body_fn, (z0, iter_count, mask))
    
    _, iterations, _ = final_state
    return iterations

# Define the grid of complex numbers
x = torch.linspace(-2.0, 1.0, 500)
y = torch.linspace(-1.5, 1.5, 500)
xx, yy = torch.meshgrid(x, y)
complex_grid = xx + 1j * yy

# Compute the Mandelbrot set
max_iter = 100
threshold = 2.0
mandelbrot_set = mandelbrot(complex_grid, max_iter, threshold)

# Plot the Mandelbrot set
plt.figure(figsize=(10, 10))
plt.imshow(mandelbrot_set, extent=(-2, 1, -1.5, 1.5), cmap="inferno")
plt.colorbar(label="Iteration count")
plt.title("Mandelbrot Set")
plt.xlabel("Real")
plt.ylabel("Imaginary")
plt.show()

    In general, the only non-obvious thing about while_loop is that the cond_fn is returning a tensor, not a bool, so make sure you are getting your types right, and that the shapes must be consistent from loop to loop. If you need more accumulating type behavior, look at scan!