Context Parallelism

[oliverkinch]

Implements CP for non MoE models. Implementing CP for MoEs will be in a separate PR.

Fix #31.
#38 will be redundant given this PR.

[oliverkinch]

From train.py we have

if parallel_dims.cp_enabled: # the following is necessary for CP w/ flex attention
    from torch.distributed.tensor.experimental._attention import _set_cp_global_var, _DispatchMode, _cp_options

    # set_rotate_method("alltoall")  # alltoall or allgather (only allgather for flex)
    _set_cp_global_var("cp_shard_dim", 2)
    # _cp_options.enable_load_balance = True  # no load balancing for flex
    torch.distributed.tensor.experimental._attention._dispatch_mode = (
        _DispatchMode.TORCH_FUNCTION
    )

_set_cp_global_var is only available in torch 2.9.0, but if I force this version the code crashes as .backward() is called. Is _set_cp_global_var necessary?

[oliverkinch]

Problems with FLASH_ATTENTION? It works with MATH

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    return F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=scale)

torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in function scaled_dot_product_attention>(*(FakeTensor(..., device='cuda:0', size=(1, 16, 4096, 192), dtype=torch.bfloat16,
           grad_fn=<TransposeBackward0>), FakeTensor(..., device='cuda:0', size=(1, 16, 4096, 192), dtype=torch.bfloat16,
           grad_fn=<TransposeBackward0>), FakeTensor(..., device='cuda:0', size=(1, 16, 4096, 128), dtype=torch.bfloat16,
           grad_fn=<TransposeBackward0>)), **{'is_causal': True, 'scale': 0.07216878364870322}): got RuntimeError('No available kernel. Aborting execution.')

[oliverkinch]

@rlrs Context parallelism now runs for gemma and llama

[oliverkinch]

I have now also included the related to YaRN in this PR, see d36078d

• Blue: The base model.
• Orange: The base model with its context window extended from 4k to 32k using YaRN, without any additional training.
• Green: The same YaRN-extended model, further trained for 1,000 steps on long-context data (wiki_expanded)

[oliverkinch]

All the CMDs below run as expected. Anything else we should test before merging this? @rlrs

Training
python -m torch.distributed.run --standalone --nproc-per-node 2 train.py jobs/gemma-cp2
python -m torch.distributed.run --standalone --nproc-per-node 2 train.py jobs/llama-cp2

Comparing a Maester DCP checkpoint against a Hugging Face model

python compare_models.py \
  --job-config jobs/munin-32k/config.json \
  --checkpoint-dir jobs/munin-32k/checkpoints/step-1000 \
  --hf-model oliverkinch/munin-32k-step-1000 \
  --num-prompts 0 \
  --dataset data/wiki-expanded-hf \
  --dataset-samples 4 \
  --dataset-max-length 512

Giving output as:

PROMPT 0: '# Frankrig\n\nFrankrig (fransk: "France"), officielt Den Franske Republik (fransk:'
Tokenized length: 512
Logit max abs diff:  7.428315e+00
Logit mean abs diff: 1.034052e-01
HF loss:       0.775544 (ppl=2.172)
Maester loss:  0.771280 (ppl=2.163)

Are these differences acceptable?

YaRN convert

python -u scripts/convert/llama/from_dcp_yarn.py \
  jobs/munin-32k/checkpoints \
  /tmp/munin-open-7b-pt-export \
  --name step-1000 \
  --base danish-foundation-models/munin-open-7b-pt