The check if "weight" in norm: is not a valid way to verify parameter existence on an mlx.nn.Module. This will likely evaluate to False or raise a TypeError, causing the LayerNorm calculation to skip applying the weight and bias tensors. This will result in incorrect numerical output for the Qwen3-VL vision blocks.

Thanks. The current code already does exactly this: _qwen3_torch_like_layer_norm reads getattr(norm, "weight", None) and getattr(norm, "bias", None) and only applies them when present, so the affine scale and bias are honored for the Qwen3-VL vision blocks. Worth noting for the record that mlx.nn.Module subclasses dict, so even "weight" in norm would have reported parameter presence correctly here.

This manual weight decay implementation has two significant issues:

1. Precision Underflow: Calculating the scale in the parameter's native dtype (e.g., float16 or bfloat16) will cause the weight decay to be ignored. For typical values like lr=2e-4 and wd=0.01, the term lr * wd (2e-6) is smaller than the machine epsilon for float16/bfloat16 relative to 1.0, so 1.0 - 2e-6 rounds back to 1.0.
2. Unintended Parameter Promotion: If scale is calculated in float32 (to fix the precision issue), the operation parameter * scale will promote the model parameters to float32. Since these parameters (LoRA weights and norms) are explicitly excluded from the restoration logic in _restore_trainable_storage_dtypes, they will remain in float32, doubling their memory footprint for the rest of the training session.

The calculation should be done in float32 and explicitly cast back to the original dtype.

Good catch on both points. The current _apply_manual_weight_decay computes the scale in fp32 and casts the result back to the parameter dtype: (parameter.astype(mx.float32) * scale).astype(parameter.dtype), so there is no float16/bfloat16 underflow and no silent promotion to fp32. This matches the suggested fix.

Removing the NOT_RETURN_LOGITS check here, combined with the deletion of the specialized branch for returned logits, introduces a performance and memory regression. When UNSLOTH_RETURN_LOGITS=1 is active, logits are materialized early; forcing the code into the unsloth_fused_ce_loss branch causes a redundant matmul. Furthermore, ensure that logit scaling (multiply/divide) and softcapping transformations are applied to these materialized logits before they are passed to the loss function to maintain consistency with standard execution paths.

The NOT_RETURN_LOGITS guard is still in place, so fused CE only runs when we are not returning logits. There is now a dedicated UNSLOTH_RETURN_LOGITS=1 branch that reuses the already-materialized logits and applies the logit scale-multiply, scale-divide and softcap transforms before calling the loss function, so we avoid the duplicate matmul and keep transform parity.

Exclude pad token ID from global VLM ignore list

Adding pad_token_id to _get_vlm_ignore_token_ids and then masking via _mask_label_token_ids masks that token everywhere, not only in padded positions. For tokenizers where pad_token_id == eos_token_id (common in Llama-derived setups), valid in-sequence EOS targets are turned into -100, so EOS supervision is silently removed across training/eval batches. Padding should be masked only through attention_mask == 0, not by globally ignoring the pad ID.

Useful? React with 👍 / 👎.

_get_vlm_ignore_token_ids only collects image/video/audio/boi/eoi token ids; it never adds pad_token_id to the global ignore list. Padding is masked separately via attention_mask == 0, so in-sequence EOS targets are preserved even when pad_token_id == eos_token_id.

Keep return-logits mode out of fused CE loss branch

This condition now routes all labeled ForCausalLMLoss calls through unsloth_fused_ce_loss, including UNSLOTH_RETURN_LOGITS=1. In that mode the compiler replacement already materializes logits up front, so this branch computes a second LM-head projection inside fused CE, increasing memory/compute and risking OOM on large models. The previous NOT_RETURN_LOGITS guard avoided that duplicate matmul.

Useful? React with 👍 / 👎.

The duplicate matmul concern no longer applies. The fused CE branch is gated by NOT_RETURN_LOGITS, and the UNSLOTH_RETURN_LOGITS=1 path is handled separately by reusing the materialized logits rather than re-projecting inside fused CE.

Apply LayerNorm affine params in Qwen3 vision block patch

The new _qwen3_torch_like_layer_norm gate checks use membership ("weight" in norm, "bias" in norm) instead of attribute existence. For module-style LayerNorm objects, this test does not reliably indicate parameter presence, so the affine scale/bias can be skipped (or error), making every patched Qwen3-VL vision block run a non-affine normalization that diverges from pretrained/Transformers behavior and can degrade training and loss parity.

Useful? React with 👍 / 👎.

The current helper uses getattr(norm, "weight", None) and getattr(norm, "bias", None), so the affine parameters are applied. For reference, mlx.nn.Module subclasses dict, so the membership form also reflects parameter presence exactly (LayerNorm with affine has both weight and bias; without affine has neither), so this path runs the correct affine LayerNorm.

Normalize nullable seeds before VLM epoch reshuffle

When create_vlm_batches rolls to a new epoch, it unconditionally computes int(seed) + epoch for both torch_randperm and default reshuffling. If callers pass seed=None (which is already supported in _torch_randperm_order and by NumPy seeding), multi-epoch materialization will raise TypeError on the first rollover and stop training/eval batch creation.

Useful? React with 👍 / 👎.

create_vlm_batches now normalizes the seed up front with base_seed = _normalize_seed(seed) (None defaults to 3407) before doing any base_seed + epoch arithmetic, so a seed=None caller no longer raises on epoch rollover.

Normalize nullable seed in streaming torch_randperm path

When dataset_order="torch_randperm", the streaming iterator now does int(seed) + epoch before calling _torch_randperm_order. If callers pass seed=None (which _torch_randperm_order explicitly supports by defaulting to 3407), this raises TypeError on the first epoch and stops batch generation. Please normalize seed the same way as _torch_randperm_order (or pass it through directly) before adding epoch.

Useful? React with 👍 / 👎.

Same fix as the VLM path: the streaming and ordered text paths normalize the seed with _normalize_seed before + epoch, matching _torch_randperm_order having its own None default, so seed=None is handled.

Handle scalar learning rates in manual AdamW decay

_apply_manual_adamw_weight_decay assumes optimizer.learning_rate has an .astype(...) method, but in valid configurations it can be a Python float (notably the constant scheduler path without warmup, and the MLX simulation optimizer stubs). In that case this line raises AttributeError before any optimizer step, so AdamW training fails as soon as manual decay runs. Please normalize LR to an mx.array (or branch on scalar) before dtype-casting.

Useful? React with 👍 / 👎.

_apply_manual_weight_decay now branches on whether the learning rate exposes .astype: scalar floats are wrapped via mx.array(lr_value, dtype=mx.float32), so the constant-schedule and simulation-optimizer paths no longer hit an AttributeError.

Raise on empty ordered token stream instead of returning []

When create_ordered_batches drops all rows via if len(ids) >= 2 (for example very small max_seq_length or single-token rows), it returns an empty batch list. The new ordered-data path in MLXTrainer then indexes batches with batch_idx % len(batches), which crashes at runtime with division-by-zero instead of surfacing a data error. This should raise a clear ValueError here, matching the other dataset-prep guards.

Useful? React with 👍 / 👎.

create_ordered_batches now raises a clear ValueError when no trainable sequences survive (need at least two tokens after formatting and truncation), so we surface a data error instead of producing an empty list that later divides by zero.

If this is specific to qwen3 do we need the if check?
or even in qwen3 family some norms have weight/bias and some dont?

Yes, the check is needed, and yes, even within the qwen3 family the norms differ. In mlx-vlm qwen3_vl the vision blocks use nn.LayerNorm (weight and bias both present) while the language stack uses nn.RMSNorm (weight only, no bias). Because mlx.nn.Module subclasses dict, "weight" in norm / "bias" in norm report exactly which parameters exist, so the conditional lets the same helper apply bias only where it exists and skip it for RMSNorm. The current code uses getattr(norm, "weight", None) which is equivalent. Keeping the check also keeps the helper correct if an affine=False norm ever appears.

I'm thinking we do this to avoid 0 division. But div by 1e-12 might cause the numbers to blow up?
For a similar thing what I did was set the scale to 1

Agreed, and the current code does exactly what you suggested. The NF4 dequant uses denom = mx.where(absmax > 0, absmax, mx.ones_like(absmax)) (and the same for the nested-scale denom), so zero scales divide by 1 instead of a tiny epsilon. There is no 1e-12 left anywhere in the MLX path, so there is no blow-up risk.

Keep torch_randperm active across epoch-based text training

When dataset_order="torch_randperm" is used with epoch-based training (max_steps == 0), this early break materializes only one epoch of randomized order and returns. MLXTrainer then reuses that fixed batch list with batch_idx % len(batches) for subsequent epochs, so later epochs never get a fresh randperm order. This diverges from the expected per-epoch reshuffle behavior of a torch-style random sampler and can bias training/eval results for multi-epoch runs.

Useful? React with 👍 / 👎.

The text ordered path now reseeds per epoch: make_order(epoch) calls _torch_randperm_order(len, base_seed + epoch) and is re-invoked at each epoch boundary, with the loop bounded by num_epochs * len(dataset), so multi-epoch runs get a fresh permutation each epoch.

Rebuild VLM torch_randperm order for each epoch run

In pre-materialized VLM batching, this branch exits as soon as one dataset pass completes whenever num_batches is None (the normal epoch-based path). With dataset_order="torch_randperm", that means only the first epoch gets randomized; later epochs in MLXTrainer reuse the same stored batches instead of generating a new permutation. This breaks sampler parity for multi-epoch VLM training and can skew convergence behavior.

Useful? React with 👍 / 👎.

create_vlm_batches rebuilds the order on each epoch rollover with _torch_randperm_order(len, base_seed + epoch), so later epochs no longer reuse the first epoch permutation.

Stop batches from spanning torch_randperm epoch boundaries

create_ordered_batches wraps to a new permutation inside the same batch (order_pos >= len(order)), so when num_epochs is used and len(dataset) % batch_size != 0, one batch can contain samples from two different epochs. This changes both sample order and batch count (e.g., 5 samples, batch size 4, 2 epochs yields 3 batches instead of the expected 4), so MLXTrainer under-trains in epoch mode and no longer mirrors CUDA/DataLoader epoch semantics for dataset_order="torch_randperm".

Useful? React with 👍 / 👎.

create_ordered_batches now breaks the batch at the epoch boundary: when order_pos >= len(order) it emits the partial batch and restarts at epoch+1 rather than wrapping inside one batch, so no batch spans two epochs and the batch count matches CUDA DataLoader epoch semantics.

Exclude ln_ normalization weights from AdamW decay*

The new HF-style decay filter only checks for path segments containing "norm", so normalization parameters named like ln_1.weight or ln_f.weight still get decayed. This repo already treats ln_* as norm parameters in _ensure_lora_frozen (see _NORM_FRAGMENTS), so models using those names will now receive unintended weight decay on norm scales, diverging from the stated “skip bias and norms” behavior and changing optimization dynamics.

Useful? React with 👍 / 👎.

_should_apply_weight_decay uses _part_is_norm, which matches "norm", ln_*, and ln_f, so ln_1.weight / ln_f.weight style norm scales are excluded from decay as intended.

Include ln_ params in fp32 norm-parameter preservation*

_keep_norm_parameters_float32 claims to keep normalization parameters in fp32, but _is_norm_parameter_path only matches components containing "norm". Any normalization layer named with ln_* (which this codebase already recognizes as norm-like elsewhere) is skipped and left in lower precision, undermining the stabilization this pass is meant to provide for FT/LoRA/QLoRA training.

Useful? React with 👍 / 👎.

The loader _is_norm_parameter_path matches "norm", ln_*, and ln_f as well, so GPT-2 / GPT-OSS style ln_* normalization parameters are kept in fp32 by _keep_norm_parameters_float32.

Preserve the prebuilt epoch flag before step counting

In completion-only text training (train_on_responses_only) with num_train_epochs > 1 and max_steps <= 0, _create_labeled_batches already materializes all epochs and sets trainer._prepared_batches_include_epochs = True, but this reset clears that flag before _prepare_data() returns the prebuilt self._batches. The total-step calculation then treats len(batches) as one epoch and multiplies by num_train_epochs again, so a 3-epoch run executes 9 epochs' worth of micro-batches and over-trains/reuses data unexpectedly.

Useful? React with 👍 / 👎.

Confirmed and fixed. _train_inner was zeroing _prepared_batches_include_epochs before _prepare_data returned the prebuilt batches from train_on_responses_only, and since the prebuilt path returns early it never re-derived the flag, so the step counter multiplied the already-materialized epochs again. The reset is now gated on self._batches is None, which preserves the flag for the completion-only path while normal runs still reset and re-derive it. Added a regression test pinning 18 steps for a 3-epoch run instead of 54. Fixed in c8ad7dd.

Normalize nullable seed before labeled randperm ordering

When train_on_responses_only is used with dataset_order="torch_randperm" and a caller leaves seed=None, this addition raises TypeError before _torch_randperm_order can apply its own None default. The unlabeled ordered-batch path already normalizes nullable seeds, so completion-only text training is the path that still fails for the same valid seed configuration.

Useful? React with 👍 / 👎.

Good catch, fixed. The labeled torch_randperm path did seed + epoch_idx directly, which raised TypeError when seed was None, whereas the unlabeled path already runs _normalize_seed first. The seed is now normalized before the per-epoch add so seed=None falls back to the default and per-epoch reseeding still works. Added tests that build labeled batches with seed=None and assert each epoch is a distinct permutation. Fixed in c8ad7dd.

Read Gemma sliding-window settings from config

For Gemma-style VLM language stacks that pass token_type_ids, the sliding-window settings live on stack.config (sliding_window_pattern / sliding_window) rather than as direct stack attributes. These getattr calls therefore resolve to None, so every layer is treated as global and uses the full token-type mask instead of the local sliding-window mask, changing Gemma3 attention semantics and training parity whenever this custom hidden-stack path is used.

Useful? React with 👍 / 👎.

Good catch on the robustness angle. On real mlx-vlm 0.6.2 both Gemma3Model and Gemma4TextModel copy window_size and sliding_window_pattern onto the module, so today the masks are correct. A config fallback was added anyway (reading sliding_window and sliding_window_pattern off stack.config) so a stack that only stores them there still builds windowed masks rather than treating every layer as global. Fixed in c8ad7dd.