fix: correct attention class retrieval for gemma3n model in lora_kernels.py

* Support --lora-on-cpu flag for DPO model merging

* fix: use device=cpu in _convert_embedding_modules_dtype when lora_on_cpu is set

src/axolotl/common/datasets.py

@@ -75,13 +75,17 @@ def load_datasets(
 
         num_examples = cli_args.debug_num_examples if cli_args else 1
         text_only = cli_args.debug_text_only if cli_args else False
-        train_samples = sample_dataset(train_dataset, num_examples)
+        try:
-        check_dataset_labels(
+            train_samples = sample_dataset(train_dataset, num_examples)
-            train_samples,
+            check_dataset_labels(
-            tokenizer,
+                train_samples,
-            num_examples=num_examples,
+                tokenizer,
-            text_only=text_only,
+                num_examples=num_examples,
-        )
+                text_only=text_only,
+            )
+        except AttributeError:
+            # can't sample iterable datasets
+            pass
 
         LOG.info("printing prompters...")
         for prompter in prompters:

src/axolotl/core/builders/causal.py

@@ -413,7 +413,8 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 or self.cfg.micro_batch_size > 1
             ):
                 return DataCollatorForSeq2Seq(self.tokenizer, **kwargs)
-            return None
+            if not (self.cfg.sample_packing and self.cfg.pretrain_multipack_attn):
+                return None
 
         if self.cfg.model_config_type == "mamba":
             return MambaDataCollator(tokenizer=self.tokenizer)

src/axolotl/loaders/model.py

@@ -776,6 +776,9 @@ class ModelLoader:
         dist_dtype: torch.dtype,
         before_kbit_train_or_finetune: bool,
     ):
+        dest = {"dtype": dist_dtype}
+        if self.cfg.lora_on_cpu:
+            dest["device"] = "cpu"
         for name, module in self.model.named_modules():
             if "norm" in name:
                 module.to(dist_dtype)
@@ -786,4 +789,4 @@ class ModelLoader:
                     # don't upcast lm_head for btlm
                     continue
             if any(m in name for m in embedding_modules) and hasattr(module, "weight"):
-                module.to(dist_dtype)
+                module.to(**dest)

src/axolotl/monkeypatch/lora_kernels.py

@@ -156,8 +156,12 @@ def get_attention_cls_from_config(cfg: DictDefault) -> Type[nn.Module]:
         model_cls_prefix = "".join(
             [part.capitalize() for part in model_type.split("_")]
         )
-        module = __import__(module_path, fromlist=[f"{model_cls_prefix}Attention"])
+        if model_type == "gemma3n":
-        attention_cls = getattr(module, f"{model_cls_prefix}Attention")
+            module = __import__(module_path, fromlist=[f"{model_cls_prefix}TextAttention"])
+            attention_cls = getattr(module, f"{model_cls_prefix}TextAttention")
+        else:
+            module = __import__(module_path, fromlist=[f"{model_cls_prefix}Attention"])
+            attention_cls = getattr(module, f"{model_cls_prefix}Attention")
 
         return attention_cls
     except (ImportError, AttributeError) as e:

src/axolotl/monkeypatch/multipack.py

@@ -42,6 +42,10 @@ def patch_for_multipack(model_type, model_name=None, has_remote_code=False):
     if has_remote_code:
         patch_remote(model_name)
     elif hasattr(transformers, "modeling_flash_attention_utils"):
+        # sanity check in case upstream api changes on this
+        assert hasattr(
+            transformers.modeling_flash_attention_utils, "_get_unpad_data"
+        ), "transformers api changed for _get_unpad_data for flash attention"
         transformers.modeling_flash_attention_utils._get_unpad_data = (  # pylint: disable=protected-access
             get_unpad_data
         )

src/axolotl/prompt_strategies/chat_template.py

@@ -103,6 +103,7 @@ class ChatTemplatePrompter(Prompter):
         chat_template_kwargs = {
             "chat_template": self.chat_template,
             "add_generation_prompt": add_generation_prompt,
+            **self.chat_template_kwargs,
         }
 
         if tools:

src/axolotl/train.py

@@ -223,6 +223,8 @@ def execute_training(
             )
 
         LOG.info("Starting trainer...")
+        if cfg.bf16:
+            torch.set_default_dtype(torch.bfloat16)
         trainer.train(resume_from_checkpoint=resume_from_checkpoint)
 
 

src/axolotl/utils/data/pretraining.py

@@ -224,10 +224,10 @@ def wrap_pretraining_dataset(
     remove_columns = []
     if dataset.features is None:
         for first_row in dataset:
-            remove_columns = first_row.keys()
+            remove_columns = list(first_row.keys())
             break
     else:
-        remove_columns = dataset.features.keys()
+        remove_columns = list(dataset.features.keys())
 
     dataset = dataset.map(
         encode,
@@ -267,6 +267,7 @@ def encode_packed_pretraining(
         batch_size=1,
         batch_max_len=batch_size * max_seq_length,
         drop_last=True,
+        num_processes=1,
     )
 
     chunked_data = defaultdict(list)

src/axolotl/utils/data/sft.py

@@ -334,7 +334,10 @@ def _load_raw_datasets(
     dataset = merge_datasets(datasets, cfg)
 
     if not cfg.skip_prepare_dataset:
-        dataset = drop_long_seq_in_dataset(dataset, cfg)
+        if split == "test" and cfg.eval_sequence_len:
+            dataset = drop_long_seq_in_dataset(dataset, cfg.eval_sequence_len, cfg)
+        else:
+            dataset = drop_long_seq_in_dataset(dataset, cfg.sequence_len, cfg)
         if cfg.sample_packing:
             dataset, _ = process_datasets_for_packing(cfg, dataset, None)
 

src/axolotl/utils/data/utils.py

@@ -148,11 +148,14 @@ def deduplicate_and_log_datasets(
     return dataset, other_dataset
 
 
-def drop_long_seq_in_dataset(dataset: Dataset, cfg: DictDefault) -> Dataset:
+def drop_long_seq_in_dataset(
+    dataset: Dataset, sequence_len: int, cfg: DictDefault
+) -> Dataset:
     """Remove sequences longer than configured maximum from dataset.
 
     Args:
         dataset: Dataset to filter.
+        sequence_len: Maximum length for sequences to keep
         cfg: Dictionary mapping `axolotl` config keys to values.
 
     Returns:
@@ -167,7 +170,7 @@ def drop_long_seq_in_dataset(dataset: Dataset, cfg: DictDefault) -> Dataset:
 
     drop_long = functools.partial(
         drop_long_seq,
-        sequence_len=cfg.sequence_len,
+        sequence_len=sequence_len,
         min_sequence_len=cfg.min_sample_len,
     )
 
@@ -187,7 +190,7 @@ def drop_long_seq_in_dataset(dataset: Dataset, cfg: DictDefault) -> Dataset:
 
     drop_long_kwargs = {}
     if filter_map_kwargs:
-        drop_long_kwargs["desc"] = "Dropping Long Sequences"
+        drop_long_kwargs["desc"] = f"Dropping Long Sequences (>{sequence_len})"
 
     dataset = dataset.filter(
         drop_long,

src/axolotl/utils/samplers/multipack.py

@@ -260,7 +260,7 @@ class MultipackBatchSampler(BatchSampler):
         lengths: np.ndarray,  # Sequence lengths
         packing_efficiency_estimate: float = 1.0,  # Initial efficiency estimate
         drop_last: bool = True,  # Whether to drop final batches (might be incomplete)
-        num_count_samples: int = 8,  # Number of times to estimate batch count
+        num_count_samples: int = 4,  # Number of times to estimate batch count
         sequential: bool = False,  # Whether to use sequential packing
         group_size: int = 100_000,  # Size of groups for parallel packing
         bin_size: int = 200,  # The max number of samples that can be packed in a single bin
@@ -335,12 +335,13 @@ class MultipackBatchSampler(BatchSampler):
             bins = [[indices[b_idx] for b_idx in bin_indices] for bin_indices in bins]
         else:
             # Use parallel packing
+            num_processes = self.num_processes or 1
             all_bins = pack_parallel(
                 lengths,
                 bin_capacity=self.batch_max_len,
                 group_size=self.group_size,
                 bin_size=self.bin_size,
-                num_processes=max(4, self.num_processes) if self.num_processes else 4,
+                num_processes=min(4, num_processes) if num_processes else 4,
                 safe_mode=self.safe_mode,
                 mp_start_method=self.mp_start_method,
             )

src/axolotl/utils/schemas/config.py

@@ -366,6 +366,12 @@ class AxolotlInputConfig(
             "description": "The maximum length of an input to train with, this should typically be less than 2048 as most models have a token/context limit of 2048"
         },
     )
+    eval_sequence_len: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The maximum length of an input for evaluation. If not specified, defaults to sequence_len"
+        },
+    )
     min_sample_len: int | None = None
     max_prompt_len: int = Field(
         default=512,
@@ -778,6 +784,12 @@ class AxolotlInputConfig(
             "description": "Custom jinja template for chat template. This will be only used if chat_template is set to `jinja` or `null` (in which case chat_template is automatically set to `jinja`). Default is null."
         },
     )
+    chat_template_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Additional kwargs to pass to the chat template. This is useful for customizing the chat template. For example, you can pass `thinking=False` to add a generation prompt to the chat template."
+        },
+    )
     eot_tokens: list[str] | None = Field(
         default=None,
         json_schema_extra={

src/axolotl/utils/schemas/validation.py

@@ -462,6 +462,20 @@ class TrainingValidationMixin:
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def pretrain_with_tps(cls, data):
+        if data.get("pretraining_dataset") and data.get(
+            "include_tokens_per_second", False
+        ):
+            # combining these would raise `TypeError: cannot pickle 'dict_keys' object`
+            # due to trying to count the number of tokens total in the dataset
+            raise ValueError(
+                "pretraining_dataset and include_tokens_per_second cannot be used together."
+            )
+
+        return data
+
 
 class LoRAValidationMixin:
     """Validation methods related to LoRA/QLoRA configuration."""

src/axolotl/utils/trainer.py

@@ -381,6 +381,7 @@ def process_pretraining_datasets_for_packing(
     if not skip_position_ids:
         train_dataset = train_dataset.map(
             add_position_ids,
+            batched=True,
             desc="Add position_id column (Pretraining Sample Packing)",
         )
     if drop_attention_mask:

tests/test_packed_batch_sampler.py

@@ -70,7 +70,7 @@ class TestBatchedSamplerPacking:
         )
         train_dataset = concatenate_datasets([dataset_wrapper])
 
-        train_dataset = drop_long_seq_in_dataset(train_dataset, cfg)
+        train_dataset = drop_long_seq_in_dataset(train_dataset, cfg.sequence_len, cfg)
 
         lengths = get_dataset_lengths(train_dataset)
         batch_sampler = MultipackBatchSampler(