Претрейн char-level языковых моделей

За основу взят код run_clm.py из библиотеки transformers.

В него внесены некоторые модификации:

1. Добавлено использование charlevel_tokenizer - посимвольного токенизатора, поддерживающего необходимый API токенизаторов transformers. Чтобы его использовать, надо указать опцию

--tokenizer_name путь_к_директории_с_файлом_tokenizer_config.json

Код run_clm.py увидит, что в конфиге фигурирует CharacterTokenizer, и создаст соответствующий экземпляр вместо AutoTokenizer.from_pretrained().

2. Поправлена работа с IterableDataset: когда задана опция --streaming true, выполняется перемешивание фрагментов датасета:

raw_datasets = raw_datasets.shuffle(buffer_size=5_000_000, seed=42)

Благодаря этому уменьшается влияние детерминированной последовательности шардов датасета на процесс обучения. Это важно, если корпус претрейна формируется последовательно из множества тематических источников, например сначала идет Википедия, потом большой массив худлита, затем новости и т.д. Размер буфера для перемешивания (параметр buffer_size) может быть неоптимален для вашего датасета - обратите на это внимание.

3. Немного улучшена работа с датасетами претрейна, в которых сделано шардирование данных в набор json-файлов, за счет корректного учета файла конфигурации датасета. Организовано всё так.

Создается датасет, в котором данные претрейна раскиданы на множество json-файлов такого формата:

{
    "texts": [
        {
            "text": "строка длиной N символов"
        },
        {
            "text": "строка длиной N символов"
        },
        ...
    ]

Каждая запись в таком файле имеет фиксированную длину, соответствующую размеру окна модели.

Кроме шардов с данными, в директории датасета есть файл конфига примерно такого вида:

{
    "field": "texts",
    "block_size": 1024,
    "features": {
        "text": "string"
    }
}

В нем параметр "block_size" равен N в шардах.

Далее, с помощью опции --dataset_config_name путь_к_конфигу_датасета указываем путь к этому конфигу. Код в run_clm.py загрузит конфиг и настроит итератор по данным соответствующим образом. По большому счету, работа с конфигурацией датасета использует стандартную функциональность библиотеки datasets, нужно было добавить параметр при вызове datasets.load_dataset(...).

Пример полного скрипта запуска претрейна: pretrain_charllama.sh. Примерно так я тренировал модельку charllama-35M.

Претрейн charmamba

Вариант кода для претрейна модели с архитектурой Mamba лежит отдельно, так как требует, чтобы был установлен пакет из форка mamba.