YaLM 100B Documentation

This implementation provides an example of how to use DeepSpeed's curriculum learning (CL) feature to achieve faster and more stable language model pre-training.

Key constraints and details:

• Supported Architectures: Currently, this is only integrated for GPT pre-training.
• Implementation Context: This is one of two curriculum learning examples for Megatron-LM GPT-2 pre-training. Users should be aware that different implementations may have unique features and limitations.
• External Resources:
  • For a detailed guide, refer to the DeepSpeed curriculum learning tutorial.
  • For the underlying technical methodology, refer to the original paper.

GPT-2 data preprocessing requires specific modifications compared to BERT: you must provide a merge table, use the GPT2BPETokenizer, append an end-of-document token (--append-eod), and do not split sentences.

The output files will be named _text_document.bin and _text_document.idx.

python tools/preprocess_data.py \
       --input my-corpus.json \
       --output-prefix my-gpt2 \
       --vocab gpt2-vocab.json \
       --dataset-impl mmap \
       --tokenizer-type GPT2BPETokenizer \
       --merge-file gpt2-merges.txt \
       --append-eod

To prepare the training dataset, you must install several Python libraries and the LSH (Locality Sensitive Hashing) dependency. Run the following commands to set up the environment:

    pip install ftfy langdetect numpy torch pandas nltk sentencepiece boto3 tqdm regex bs4 newspaper3k htmlmin tldextract 
    git clone https://github.com/mattilyra/LSH
    cd LSH
    python setup.py install

The repository provides several shell scripts for different generation modes:

• Interactive Generation: Use examples/generate_interactive.sh for simple command-line interaction.
• Conditional Sampling: Use examples/generate_conditional_sampling.sh. This uses a sampling strategy (Top-p by default, but temperature and top-k can be adjusted). Input must be in jsonlines format (e.g., examples/example_cond_input.json). The output will be jsonlines with a generated text field added to each line.
• Conditional Greedy: Use examples/generate_conditional_greedy.sh. This is suitable for few-shot problem solving.
• Unconditional Generation: Use examples/generate_unconditional.sh. No input is required; the output is provided in jsonlines format.

To fine-tune a model from a pretrained checkpoint on a new corpus, add the --finetune flag to your training script and adjust the input files and training parameters.

Warning: Adding --finetune resets the iteration count to zero and reinitializes the optimizer and internal state. If a fine-tuning session is interrupted, you must remove the --finetune flag before resuming, otherwise training will restart from the beginning.

Once an ICT model is trained, you can build a BlockData structure to enable fast similarity search using FaissMIPSIndex. Use tools/create_doc_index.py to embed the entire dataset. This script can be run in an interactive session and supports multi-GPU/multi-node execution for large datasets.

python tools/create_doc_index.py \
    --num-layers 12 \
    --hidden-size 768 \
    --ict-head-size 128 \
    --num-attention-heads 12 \
    --batch-size 128 \
    --checkpoint-activations \
    --seq-length 256 \
    --max-position-embeddings 256 \
    --ict-load /path/to/pretrained_ict \
    --data-path /path/to/indexed_dataset \
    --titles-data-path /path/to/titles_indexed_dataset \
    --block-data-path embedded_blocks.pkl \
    --indexer-log-interval 1000 \
    --indexer-batch-size 128 \
    --vocab-file /path/to/vocab.txt \
    --num-workers 2 \
    --fp16

You can download the model weights and vocabulary using the provided script or by cloning the Hugging Face repository.

Option 1: Using the download script Run the bash script to download files to the default local directories.

• Weights: ./yalm100b_checkpoint/weights/
• Vocabulary: ./yalm100b_checkpoint/vocab/

Option 2: Hugging Face Clone the yandex/yalm-100b repository and pull the yalm100b_checkpoint directory.

bash download/download.sh

After obtaining the raw JSONL file, follow this pipeline to clean, deduplicate, and shuffle the data for training:

1. Cleanup: Perform ftfy, English language detection, and remove documents with fewer than 128 tokens. This step can be sharded.
2. Find Duplicates: Use LSH to find possible duplicates. This step cannot be sharded and typically takes 12-24 hours for OpenWebText.
3. Group Similar URLs: Group URLs that are similar based on the is_similar function (default threshold: 0.9).
4. Remove Duplicates: Remove the similar documents identified in the grouping step.
5. Shuffle: Shuffle the final dataset.

You can install Megatron-LM either by cloning the repository or via pip. Ensure that python3-dev is installed on your system.

Requirements:

• Python 3.6
• PyTorch 1.5 (with GPU support)
• CUDA 10
• NCCL 2.6 or above
• NVIDIA APEX
• NLTK (required for data preprocessing, but not for training or evaluation)

Recommended Environment: Using an NVIDIA NGC PyTorch container is strongly recommended (e.g., nvcr.io/nvidia/pytorch:20.03-py3).

pip install megatron-lm