This works only for those tokenizers that only have a BOS but not a EOS token.
For those that come with both, can we control whether tokenize adds the BOS and EOS tokens independently? I'm worried that we are adding the EOS token at the end of the first segment and the BOS token at the beginning of the last segment.

good catch! I'll make it explicitly add BOS only for the first segment
Btw, most tokenizers (Llama-3.1, Mistral-Nemo-Base-2407, OLMoE-1B-7B-0924) do not add the EOS token with add_special_tokens=True. Does this mean we've been training the models without the EOS token?

In the future I think we should make this config driven. The default behaviour would be to add both BOS and EOS tokens. It's important for pretraining with attention mask, and especially for SFT.

Does this mean we've been training the models without the EOS token?

Indeed, we decided that adding both BOS and EOS tokens in pretraining was unnecessary, because they are redundant. Here though I think we need to add them because we need to teach the model to terminate a response with the EOS token so that generation can stop at the right moment. Btw, I think HF is not adding the EOS token by default because otherwise prompts would end with it.

does the triton implementation support masking?

I think it doesn't: https://github.com/ServiceNow/Fast-LLM/blob/soham/loss-masking-spans/fast_llm/functional/triton/cross_entropy.py
We need to add it. Since this is the same for all loss functions, it would make sense to implement it before dispatching to specialized cross-entropy implementations:

def cross_entropy_forward_backward(
    logits,
    target,
    grad_output: float | None,
    group: ProcessGroup | None,
    implementation: CrossEntropyImpl = CrossEntropyImpl.fused,
    logits_scale_factor: float = 1.0,
    ignore_index: int=-100,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    ...
    mask = target != ignore_index
    target = target[mask]
    logits = logits[mask]
    ...