FSDP в Accelerate для крупноконтекстного обучения LLaMA ⇐ Python

[Anonymous]

Я пытаюсь обучить модель LLaMA с большими контекстами, используя Hugging Face Trainer, Fully Sharded Data Parallel (FSDP) и библиотеку ускорения для обработки ограничений памяти. Размер моего контекста очень велик: максимальный размер токена превышает 70 тысяч. Хотя моя настройка работает нормально для контекстов меньшего размера, с большими размерами у меня возникают проблемы.
Вот упрощенная версия моей настройки:

Код: Выделить всё

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import Dataset

model_name = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=False)
tokenizer.pad_token = tokenizer.eos_token

dataset = load_dataset("my_private_dataset")

def tokenize_function(examples):
tokenized = tokenizer(
prompts,
padding="max_length",
truncation=True,
max_length=70000
)
tokenized['labels'] = tokenized['input_ids']
return tokenized

tokenized_dataset = dataset.map(tokenize_function, batched=True)

training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=5,
per_device_train_batch_size=1,
save_steps=5000,
save_strategy="epoch",
save_total_limit=1,
logging_dir='./logs',
logging_steps=1000,
eval_strategy="no",
gradient_checkpointing=True,
)

trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
)

trainer.train()

Когда я пытаюсь тренироваться с ускорением для более крупных контекстов, я получаю следующую ошибку:

RuntimeError: Размер тензор a (0) должен соответствовать размеру тензора b (2048) в неодноэлементном измерении 2

Я проверил размеры пакетов и формы тензоров для меньших входы и один графический процессор, и все выглядит правильно. Ошибка появляется только тогда, когда размер ввода становится большим.
Вопросы:

• Кто-нибудь сталкивался конкретно с этим типом ошибки с FSDP или ускорением для больших контекстов?
• Существуют ли дополнительные настройки ускорения или шаги отладки для чрезвычайно больших последовательностей токенов, которые могут помочь выявить проблему?
• Подходит ли FSDP для большого контекста?

Подробнее здесь: https://stackoverflow.com/questions/791 ... a-training