Activation checkpointing needs to keep track of the state of the random number generator, which fails with torch.compile(). Rumor has it that the latest torch nightly has this fixed, so we should try this.

❓ The question

I tried implementing grouped query attention in this pull request, but seems that pytorch's scaled_dot_product_attention doesn't support the kind of broadcasting we'd need for this. Revisit if/when this gets fixed on pytorch's end.

There's a PyTorch implementation here: https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py

The problem we have with LayerNorm (or at least PyTorch's built-in implementation of LayerNorm) is that torch's autocast behavior is hardcoded to upcast inputs to the function torch.layer_norm, and as a result using low precision LayerNorm requires some hacking that fails to work when using torch.compile(). Low precision LN gives a huge speedup with non-compiled models, so we're optimistic that getting this to work with compiled models would provide a similar speedup.

Switching to a different norm implementation could sidestep this issue, and RMSNorm is a simpler operation so we might get an even bigger speedup.

This should be pretty easy: just add a configuration field for eval batch size and what not, then initialize an eval dataloader in the same way that we initialize the train data loader and pass it to the Trainer.

The Palm paper says it improves throughput, and doesn't slow learning, at least not at large scales.

• Abhi: "Using better FDSP configuration (BF16 all-gather, limiting all gathers, and a new non-reentrant version of activation checkpointing)."
  I think this means setting mixed_precision=PURE, limit_all_gathers=true, and activation_checkpointing_reentrant=false in the fsdp_config. Added in #61
• Using Triton-based attention rather than HazyResearch Flash. We're sort of doing this already by using PyTorch's new built-in implementation.
• Low precision LayerNorm. Implemented in #61
• Use a fused version of CELoss. I think we're already doing this?
• Increase microbatch size as high as possible.

🚀 The feature, motivation and pitch

Yesterday we spoke about where responsibility for data order lives between the llm-model and llm-data workstreams. I thought it might be good to start an issue here where we can figure out who and how we should ensure that future work can reproduce our exact pretraining data order.

Use cases:

• It should be clear exactly what tokens a given checkpoint has trained on, so we can ask if a model has seen a specific thing and how long ago it has seen that thing.
• It should be possible to recreate the same order of documents using a different tokenizer to allow other research to compare to ours.

Proposed features:

• Track which document IDs are in each concatenated example and which examples are in each batch
• Have dataloaders that queue up data across devices and nodes such that sequence of document IDs is always the same, even if the boundaries between batches change

Proposed method:

• When using j * k dataloaders, dl_i , across j nodes and k devices, each data loader should sample from n documents of the pretraining corpus, doc_i, with a stride of j*k
• For example if j=1 and k=2 , dl_0[0] == doc_0 and dl_1[0] == doc_1 while dl_0[1] == doc_2 and dl_1[1] == doc_3
• Finally during training, the model will log the doc ids and associated spans in each concatenated example, and examples will be logged per batch. A rough example might be, {batch_0 : [ [ (doc_0, (start_tok_idx, end_tok_idx)), ...], [(doc_42, (start_tok_idx, end_tok_idx)),...]]}

Let me know what I can do to help support this and coordinate responsibility between llm-model and llm-data.

Alternatives

A simple though reduced alternative is to just record the tokens that are trained on without any further changes.

While this supports the ability to see what a given checkpoint has seen and when it has seen it, this may make it difficult to recover batch and document boundaries or to reproduce this data order in a different tokenizer.

Also if we don’t ensure that data order is invariant to number of nodes and devices, it we may accidentally produce different training orders across different runs. Additionally it may be difficult for people using our codebase to reproduce our data order if they don’t have the same number of nodes and devices.

Additional context

No response

There are some checkpoints that we want to keep forever, because they are part of our output. The checkpoint saving code needs to know about those.

As of the latest composer release, there's now composer.callbacks.SpeedMonitor. See mosaicml/examples#238 for example usage. It looks like the version we have (which we took directly from Mosaic's example), was copied into composer line-for-line. So the update should be easy.

We have a separate Google Cloud project for OLMo now, we should use that. We should call the bucket something like gs://olmo-checkpoints instead of gs://allennlp-olmo.

It's supposed to be faster, and at this point, it's the only way we can get the fast interconnect working in LUMI.

Mosaic tells us we need to apply this to composer to make it work: dskhudia/composer@4d55bdb

We should expect a speed-up too.

Mosaic has been preferring this over AdamW and it uses less memory since it only keeps track of momentum. We can find an implementation in this PR. Eventually this will be added to composer, but in the meantime we can just copy it over.

One experiment is, let's just keep running the 7B and see if it recovers from the spikes.

Turns out we never tried it. If this fails, we need to talk to AMD ASAP. We cannot get to 70B without this.

Can just copy over the implementation from PaLM-pytorch.

https://github.com/allenai/LLM/blob/2118db56095157474fe1c69c1702db08af2d4f74/scripts/train.py#L187

I think having a checkpoint before any training happens would be quite useful.

We promised a 7B model. Also, we need an appropriately sized model to test multi-node training at smaller scales.

Maybe wandb will take care of this for us? I opened a ticket with them.

It is suspicious that we had two slightly different models (one with biases, one without), that both spiked at exactly the same moment. This suggests there might be a data issue.

🚀 The feature, motivation and pitch

E.g. we have an LM loss at 1/2, 1/4, 1/8 of the model.

To avoid the degenerate scenario where, for example, the first half of the model does well but the 2nd half just becomes a bunch of ~identity layers, we could have a separate LM head for each loss, i.e. at each layer in the model where we apply an LM loss. So we'd have to untie the embedding matrix from the LM head.

Alternatives

No response

Additional context

No response

🚀 The feature, motivation and pitch

Iz and I were talking about this today, can be under nice to have

This is from the PaLM paper:

"The standard Transformer formulation uses k attention heads, where the input vector for each timestep is linearly projected into “query”, “key”, and “value” tensors of shape [k, h], where h is the attention head size. Here, the key/value projections are shared for each head, i.e. “key” and “value” are projected to [1, h], but “query” is still projected to shape [k, h]. We have found that this has a neutral effect on model quality and training speed (Shazeer, 2019), but results in a significant cost savings at autoregressive decoding time. This is because standard multi-headed attention has low efficiency on accelerator hardware during auto-regressive decoding, because the key/value tensors are not shared between examples, and only a single token is decoded at a time."

Alternatives

No response

Additional context

No response

We can't run in a debugger anymore.

We should either wait until #61 merges, or do this on as part of that PR.

We need to grep through and change all mentions of "DOLMA" / "Dolma" / "dolma" and also rename the Python module itself.

Basically:

# Rename Python module.
mv dolma olmo
# Find and replace all mentions using `fd` and `sd` (`brew install fd sd`)
fd -H --exclude .git | xargs sd 'DOLMA' 'OLMo'
fd -H --exclude .git | xargs sd 'Dolma' 'Olmo'
fd -H --exclude .git | xargs sd 'dolma' 'olmo'

sfantao/pytorch-lumi:sles-rocm-5.5.1-python-3.10-pytorch-v2.0.1 (on Ducker hub).

Maybe we just hard-code it to 2**16.

Each time we evaluate we should evaluate on the entirety of each subset. This will remove some noise and make it easier for us to debug where the model is going wrong with RP.

Waiting on Dao-AILab/flash-attention#132 to be resolved

SwiGLU is by far the fastest compared to ReLU and GELU. Although this isn't quite a fair comparison since SwiGLU reduces the number of parameters in the model, all else equal.

See https://wandb.ai/ai2-llm/petew-activation-benchmarks

Of course, that means we need a reliable logging solution first ...

Stas' code is here: https://github.com/stas00/toolbox/blob/master/pytorch/all_reduce_bench.py

We should run this on a single node, and multiple nodes.

❓ The question

Ideally there should be no performance penalty for dropout layers when the dropout probability is 0. But if there is, we should bypass dropout when p=0.0.

I'm currently testing this here: https://wandb.ai/ai2-llm/dropout-benchmarks

There will be 4 runs:

1. 1.2b-bf16-no-dropout: uses a patched branch that bypasses all calls to dropout (except inside of the scaled_dot_product_attention function which we have no control over). This model is compiled using the default settings.
2. 1.2b-bf16-zero-dropout: uses the usual implementation without any code changes were we still call the Dropout modules even though the dropout probability is set to 0. This model is compiled using the default settings.
3. 1.2b-bf16-no-compile-no-dropout: same as 1.2b-bf16-no-dropout except this model is NOT compiled.
4. 1.2b-bf16-no-compile-zero-dropout: same as 1.2b-bf16-zero-dropout except this model is NOT compiled.

The PaLM paper has a short section of tweaks to the vanilla Transformer architecture. We should make sure we have all of those.

Exact spec still WIP, but TODOs are basically:

1. Athena query to get Titles & Abstracts from S2AG. Form JSON blob per document of the form:

{"text": "...", "paper_id": <identifier>}

• To concatenate, " ".join([title, abstract]) should be sufficient.
• Double-check whether structured abstracts preserve whitespace.

2. Athena query to get S2ORC-OA papers. Form JSON blob per document of the form:

{"text": "...", "paper_id": <identifier>}

• To concatenate, also whitespace join to linearize structured content.
• Keep everything for now, including tables, bibliographies, etc.

3. Build a blocklist of papers. For now, should just be a single file mapping paper_ids to some note or reason for removal. To start, this should be documents that are part of the test set for Catwalk evaluation, especially Pubmed/arXiv abstract generation.

RedPajama includes data sourced from ArXiv

An alternative is unArXive

S2's LaTex dumps from ArXiv are in s3://ai2-s2-scholarphi-pipeline-prod/daq/arxiv-source-data/bymonth/

@kyleclo @soldni

🚀 The feature, motivation and pitch

No response

Alternatives

No response

Additional context

Make additional layer and adapter for running code with existing down-stream evaluation tools (e.g., HELM, Catwalk).

🐛 Describe the bug

Ideally, the Layernorm layers for a sequential block are separate for attention and FFN modules. Good thing is we haven't seen too much difference because of this until now.

        q, k, v = self.att_proj(self.norm(x)).split(self.fused_dims, dim=-1)

        # Get attention scores.
        att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache)

        # Add attention scores.
        # shape: (B, T, C)
        x = x + self.dropout(att)

        # Add feed-forward projection.
        # shape: (batch_size, seq_len, d_model)
        x = x + self.dropout(self.ff_out(self.act(self.ff_proj(self.norm(x)))))

self.norm should be self.norm2 ideally in the last line

Versions

No response

This is an experiment, but probably worthwhile. With 1000 nodes, it gets very difficult to see which node did what and when, especially when they fail, if they all log to stdout, which slurs stuffs into a single log file. Papertrail is basically running syslog-as-a-service. We can send log messages there, and then look and filter with their UI. We need to configure our code to write logs to their syslog service to do this.

🚀 The feature, motivation and pitch

We need a generate method to run evaluation. Minimum requirement:

• top_k / top_p support
• temperature control
• return log probabilities.

Llama's generate method is fairly straightforward but we can't use it because of it has the wrong license.

Alternatives

• Converting the model in HuggingFace and use their generate method. This will eventually need to be done, but, since we haven't fully finalized the architecture, this is too soon.
• Finding a Apache 2.0 generation method.

Additional context

No response

While this wasn't implemented as of PyTorch 1.13.1, it appears it's going to be in the next release because it's implemented in the master branch: https://github.com/pytorch/pytorch/blob/master/torch/distributed/fsdp/api.py#L31.