Workload Title: Decoder only LM

[priyakasimbeg]

Workload

Task

Text generation.

Dataset

TBD

Model

TBD
Possible candidates include

• [preferred starting point] Nanodo
• NanoGPT
• Meta’s lingua
• Keller Jordan’s modded nanoGPT

Reference Implementation

TBD

---

ToDo

• Preliminary Benchmarking for model/dataset
• Choose Model
• Choose Dataset
• Implement data input pipeline
  • Document specific dataset version in workload-specific README
• Add model in JAX
  • Document model in workload-specific README
• Add model in PyTorch
• Target setting for new workload
• Provide sample submission (and sample tuning search space)
  • Document results of sample submission in workload-specific README (how long did it take to reach the target performance)

[priyakasimbeg]

Looks like the PyTorch workload is about 2x slower
.
JAX: 0.7136188477 seconds / step
Pytorch: 1.47774810191 seconds / step
PyTorch logs: https://gist.github.com/priyakasimbeg/c8bd57bf93583bb64c4c859840f97621

Coommand to reproduce:

torchrun --redirects 1:0,2:0,3:0 \
    --standalone \
    --nnodes=1 \
    --nproc_per_node=4 submission_runner.py --framework=pytorch --workload=lm --submission_path=algorithms/target_setting_algorithms/fineweb_edu_lm/pytorch_nadamw_target_setting.py --data_dir=/data/fineweb_edu_10B --experiment_dir=/experiment_runs --experiment_name=tests/regression_tests/adamw/original_search_trial_64_pytorch --overwrite=True --save_checkpoints=False --tuning_ruleset=external --tuning_search_space=algorithms/target_setting_algorithms/fineweb_edu_lm/tuning_search_space.jso

[runame]

Based on the logs alone you could try setting torch.set_float32_matmul_precision('high') (docs), but this is certainly not responsible for all of the slow-down (also make sure that there is no loss regression).

The graph break appears to only affect evals, could be similar to this issue. Are the evals >2x slower than in Jax?