Instruction-tuning models on super-long text

[GianlucaDeStefano]

Hi everyone,

I’m working with a custom dataset consisting of extremely long texts—each approximately 100K tokens or more—which cannot be split into smaller subtexts. Each text is paired with a question and answer, requiring the model to reason through the entire context to generate an appropriate response.

The challenge I'm facing is the amount of GPU memory required to fine-tune any model with such lengthy contexts. With access limited to A100 GPUs with up to 80GB of memory, naive parallelism techniques are insufficient.

To address this, I’m considering fine-tuning my models using tensor parallelism with MegatronLM. My goal is to fine-tune LLaMA 3.1, which has an 8K token context window that can be extended to 128K tokens using RoPE.

I’m seeking advice on how to achieve the equivalent functionality of Hugging Face’s DataCollatorForCompletionOnlyLM in MegatronLM. Specifically, I want to compute the loss only for tokens that occur after a designated 'marker' in the text so that the loss will be computed only on the 'answer' the model is tasked to generate.
How can I implement this functionality in megatronLM? is MegatronLM the right tool for my particular usecase?

[Advait251206]

What you’re trying to do is basically:

long-context instruction tuning + loss masking (answer-only supervision)

Megatron-LM can do this, but there are a few important realities and some gaps compared to Hugging Face-style data collators.

---

1. Is Megatron-LM the right tool for 100K-token instruction tuning?

Partially, but with caveats.

Megatron-LM is excellent for:

• tensor/pipeline parallelism
• very large models
• efficient long-sequence training (FlashAttention, sequence parallel, etc.)

But for your case (100K tokens + instruction masking), the constraints are:

⚠️ Key limitation

Even if RoPE is extended to 128K:

• attention cost is still O(n²) (unless you use specific sparse/linear attention variants)
• activation memory at 100K tokens is extremely large even with TP/PP

So:

• A100 80GB + full dense attention at 100K is barely feasible or infeasible depending on batch size and model size
• Megatron helps, but it does not eliminate quadratic cost

---

2. “Completion-only loss” in Megatron-LM

What HF does with DataCollatorForCompletionOnlyLM is basically:

labels = input_ids.copy()
labels[:marker_index] = -100

In Megatron-LM, the equivalent mechanism already exists

Megatron uses:

• loss_mask
• tokens + labels preprocessing stage

The core idea is:

loss is computed only where loss_mask == 1

---

3. How to implement marker-based masking

You need to modify your dataset preprocessing (NOT the model).

Step 1: define a separator token

Example:

<|user|> question text ...
<|assistant|>
answer text ...

or a custom marker like:

### RESPONSE:

---

Step 2: build loss_mask

In dataset construction (e.g. indexed_dataset.py or preprocessing script), do:

loss_mask = np.zeros(len(tokens), dtype=np.int32)

answer_start = find_marker_position(tokens)

loss_mask[answer_start:] = 1

Then store it alongside tokens.

---

Step 3: ensure dataset returns it

Megatron training expects something like:

• input_ids
• labels
• loss_mask

or in newer mcore versions:

data["loss_mask"]

---

Step 4: verify loss computation

In Megatron core (simplified):

loss = (loss_mask * cross_entropy).sum() / loss_mask.sum()

So your masking directly controls supervision.

---

4. Important: 100K token instruction tuning is the real bottleneck

Even with Megatron:

You will hit:

• activation memory explosion
• optimizer memory (if full finetune)
• attention compute wall

You likely need at least one of:

✔ Sequence parallelism

✔ Selective recomputation

✔ FlashAttention v2/v3

✔ LoRA / QLoRA (strongly recommended)

✔ Context chunking (even if logically “not split”, internally it must be)

---

5. Practical recommendation

If your constraint is:

“cannot split text into smaller subtexts”

Then in practice you still have to choose one of these internal strategies:

Option A (best): sliding window training

• chunk 8K–16K
• overlap context
• keep full document logically

Option B: hierarchical attention / retrieval augmentation

• only feed relevant segments

Option C: extreme scaling Megatron setup

• TP + PP + SP + FlashAttention
• batch size = 1
• gradient accumulation

---

6. Bottom line

• Yes, Megatron-LM can support completion-only loss via loss masking
• No, it does not have an HF-style “collator” abstraction—you implement it in preprocessing
• 100K-token full attention fine-tuning is the hardest part, not the masking

---

If I were you

I would:

• keep Megatron for scaling model training
• implement loss_mask logic in preprocessing
• strongly reconsider full 100K dense attention training unless you have multi-node H100-class setup

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206