Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Multipack simplify for Mixtral #1142

Merged
merged 8 commits into from
Jan 18, 2024
Merged

Multipack simplify for Mixtral #1142

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
77dff8d
use existing mixtal with simpler unpad multipack logic
winglian Jan 18, 2024
23d23c1
fix mistral check
winglian Jan 18, 2024
bfc95f0
ensure cu_seqlens is on correct device and add logging
winglian Jan 18, 2024
e444919
add validation for zero3 with flash attention
winglian Jan 18, 2024
29c9f71
chore: lint
winglian Jan 18, 2024
701052a
update smoke test for mixtral w multipack
winglian Jan 18, 2024
524b873
fix test
winglian Jan 18, 2024
326c92f
Merge branch 'main' into multipack-simplify
winglian Jan 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions src/axolotl/core/trainer_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from dataclasses import dataclass, field
from functools import wraps
from pathlib import Path
from typing import Optional
from typing import Optional, Type, Union

import torch
import transformers
Expand All @@ -37,6 +37,7 @@
BatchSamplerDataCollatorForSeq2Seq,
DataCollatorForSeq2Seq,
MambaDataCollator,
V2BatchSamplerDataCollatorForSeq2Seq,
)
from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
from axolotl.utils.schedulers import (
Expand Down Expand Up @@ -896,14 +897,22 @@ def build_collator(
if is_eval and training_args.eval_sample_packing:
use_batch_sampler_collator = True

collator: Type[
Union[
V2BatchSamplerDataCollatorForSeq2Seq,
BatchSamplerDataCollatorForSeq2Seq,
DataCollatorForSeq2Seq,
]
]
if use_batch_sampler_collator:
return BatchSamplerDataCollatorForSeq2Seq(
self.tokenizer,
return_tensors="pt",
**kwargs,
)
if self.cfg.model_config_type == "mixtral":
collator = V2BatchSamplerDataCollatorForSeq2Seq
else:
collator = BatchSamplerDataCollatorForSeq2Seq
else:
collator = DataCollatorForSeq2Seq

return DataCollatorForSeq2Seq(
return collator(
self.tokenizer,
return_tensors="pt",
**kwargs,
Expand Down
18 changes: 4 additions & 14 deletions src/axolotl/monkeypatch/mixtral/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,20 +3,10 @@
"""
import transformers

from axolotl.monkeypatch.utils import get_unpad_data

def replace_mixtral_attn_with_multipack_flash_attn():
from .modeling_mixtral import (
MixtralMultipackFlashAttention2,
mixtral_decoder_layer_forward,
mixtral_model_forward,
)

transformers.models.mixtral.modeling_mixtral.MixtralDecoderLayer.forward = (
mixtral_decoder_layer_forward
)
transformers.models.mixtral.modeling_mixtral.MixtralModel.forward = (
mixtral_model_forward
def replace_mixtral_attn_with_multipack_flash_attn():
transformers.models.mixtral.modeling_mixtral._get_unpad_data = ( # pylint: disable=protected-access
get_unpad_data
)
transformers.models.mixtral.modeling_mixtral.MIXTRAL_ATTENTION_CLASSES[
"flash_attention_2"
] = MixtralMultipackFlashAttention2
Loading