diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md index c20476bde45..7bc0477e705 100644 --- a/examples/post_training/modelopt/README.md +++ b/examples/post_training/modelopt/README.md @@ -54,7 +54,7 @@ to try our latest features. > be downloaded and provided through `${HF_MODEL_CKPT}`. -### ⭐ NVFP4 Quantization, Qauntization-Aware Training, and Model Export +### ⭐ NVFP4 Quantization, Quantization-Aware Training, and Model Export Provide the pretrained checkpoint path through variable `${HF_MODEL_CKPT}` and provide variable `${MLM_MODEL_SAVE}` which stores a resumeable Megatron-LM distributed checkpoint. To export @@ -97,6 +97,47 @@ export the model with flag `--export-vllm-fq`: For KV cache quantization, add a flag like `MLM_EXTRA_ARGS="--export-kv-cache-quant fp8"` while specifying your desired KV cache precision (see `KV_QUANT_CFG_CHOICES` in `quantize.py`). +### ⭐ Auto Quantize (Mixed-Precision Search) + +Auto Quantize uses `mtq.auto_quantize` to perform a per-layer mixed-precision search, assigning each +layer the best quantization format (e.g. NVFP4 or FP8) subject to a target effective-bits constraint. +This produces a model that is more accurate than uniform quantization at the same average bit-width. + +Pass `auto` as the second positional argument to `quantize.sh` and provide `--auto-quantize-bits` +through `MLM_EXTRA_ARGS`. The script will skip `--export-quant-cfg` entirely and drive the search +via the auto-quantize arguments. + +> **Note:** Auto Quantize requires `--pipeline-model-parallel-size 1` (PP=1) and +> [Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer) **0.46 or greater** +> (`pip install nvidia-modelopt>=0.46`). Alternatively, install from the +> [main branch](https://github.com/NVIDIA/Model-Optimizer) for the latest features. + +```sh +\ + TP=1 \ + HF_MODEL_CKPT= \ + MLM_MODEL_SAVE=/tmp/Llama-3.2-1B-Instruct_auto_quant \ + MLM_EXTRA_ARGS="--auto-quantize-bits 4.0" \ + ./quantize.sh meta-llama/Llama-3.2-1B-Instruct auto + +\ + PP=1 \ + HF_MODEL_CKPT= \ + MLM_MODEL_CKPT=/tmp/Llama-3.2-1B-Instruct_auto_quant \ + EXPORT_DIR=/tmp/Llama-3.2-1B-Instruct_auto_quant_export \ + ./export.sh meta-llama/Llama-3.2-1B-Instruct +``` + +Key arguments (passed via `MLM_EXTRA_ARGS`): + +| Argument | Default | Description | +| --- | --- | --- | +| `--auto-quantize-bits` | *(required)* | Target effective bits per weight (e.g. `4.0`, `4.8`). | +| `--auto-quantize-formats` | `NVFP4_DEFAULT_CFG FP8_DEFAULT_CFG` | Space-separated list of quant configs to search over. | +| `--auto-quantize-method` | `gradient` | Sensitivity scoring method (`gradient` or `kl_div`). | +| `--auto-quantize-score-size` | `128` | Number of samples used for sensitivity scoring. | +| `--auto-quantize-checkpoint` | `None` | Optional path to save/restore search state across runs. | + ### ⭐ Online BF16 EAGLE3 Training Online EAGLE3 training has both the target (frozen) and draft models in the memory where the `hidden_states` diff --git a/examples/post_training/modelopt/finetune.py b/examples/post_training/modelopt/finetune.py index 006a559aa71..912e985a820 100755 --- a/examples/post_training/modelopt/finetune.py +++ b/examples/post_training/modelopt/finetune.py @@ -16,14 +16,13 @@ from megatron.core import mpu, tensor_parallel from megatron.core.enums import ModelType from megatron.core.models.gpt import GPTModel -from megatron.core.utils import get_batch_on_this_cp_rank from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.loss_func import loss_func from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.non_loss_data_func import report_draft_acceptance_length from megatron.training import get_args, get_timers, pretrain -from megatron.training.utils import get_ltor_masks_and_position_ids, print_rank_0 -from utils import get_hf_tokenizer +from megatron.training.utils import print_rank_0 +from utils import build_lm_batch, get_eos_token_id, get_hf_tokenizer from model_provider import model_provider from megatron.core.parallel_state import get_context_parallel_group @@ -42,25 +41,6 @@ def add_finetune_args(parser): add_modelopt_args(parser) return parser -def get_eos_id(): - """Return the eos token id. - - We insert eos_token between two samples during packing. However, if the eos_token is used in message or after turns, - we need to replace it with some other special tokens that do not appear in message.""" - hf_tokenizer = get_hf_tokenizer() - - if hf_tokenizer.eos_token == "<|eot_id|>": - return 128001 - if hf_tokenizer.eos_token == "<|eot|>": - return 200001 - if hf_tokenizer.eos_token == "<|im_end|>": - return 151643 - if hf_tokenizer.eos_token == "<|return|>": - return 199999 - - return hf_tokenizer.eos_token_id - - class OfflineDataset(torch.utils.data.Dataset): def __init__(self, data_dir: str, num_samples): self.data_dir = data_dir @@ -283,7 +263,7 @@ def _process_example(self, example: Dict[str, Any]): # We always add eos between samples for training purpose. input_ids = self.tokenizer.apply_chat_template(example) current_loss_mask = [1] * len(input_ids) - input_ids = input_ids + [get_eos_id()] + input_ids = input_ids + [get_eos_token_id(self.tokenizer)] current_loss_mask += [0] assert len(input_ids) == len(current_loss_mask) @@ -396,7 +376,7 @@ def get_batch(data_iterator): datatype = torch.int64 data_b = tensor_parallel.broadcast_data(keys, data, datatype) data_b["loss_mask"] = torch.ones_like(data_b["input_ids"]) - data_b["loss_mask"][data_b["loss_mask"]==get_eos_id()] = 0 + data_b["loss_mask"][data_b["loss_mask"] == get_eos_token_id()] = 0 data_b["loss_mask"] = torch.cat([data_b["loss_mask"], torch.zeros(1,1).to(torch.cuda.current_device())], dim=-1) keys = ["aux_hidden_states", "hidden_states"] @@ -404,36 +384,21 @@ def get_batch(data_iterator): feature_b = tensor_parallel.broadcast_data(keys, data, datatype) - # Unpack the data received. - tokens_ = data_b["input_ids"] - tokens = tokens_[:, 0 : 0 + args.seq_length].contiguous() - labels = tokens_[:, 1 : 1 + args.seq_length].contiguous() - answer_only_loss_mask = data_b["loss_mask"][:, 1 : 1 + args.seq_length].contiguous() - - # Get the masks and postition ids. - attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( - tokens, get_eos_id(), get_eos_id(), args.reset_position_ids, args.reset_attention_mask, args.eod_mask_loss, False + sample_loss_mask = data_b.get("loss_mask") + batch = build_lm_batch( + data_b["input_ids"], + args.seq_length, + sample_loss_mask=sample_loss_mask, + eos_token_id=get_eos_token_id(), + reset_position_ids=args.reset_position_ids, + reset_attention_mask=args.reset_attention_mask, + eod_mask_loss=args.eod_mask_loss, + cp_group=get_context_parallel_group(), ) - loss_mask = loss_mask * answer_only_loss_mask.to(dtype=loss_mask.dtype) - - - labels = labels.contiguous() - loss_mask = loss_mask.contiguous() - - batch = { - "tokens": tokens, - "labels": labels, - "loss_mask": loss_mask, - "attention_mask": attention_mask, - "position_ids": position_ids, - } if args.export_offline_model: - batch["aux_hidden_states"] = feature_b["aux_hidden_states"].transpose(0, 1)[:args.seq_length] - batch["hidden_states"] = feature_b["hidden_states"].transpose(0, 1)[:args.seq_length] - - # slice batch along sequence dimension for context parallelism - batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=get_context_parallel_group()) + batch["aux_hidden_states"] = feature_b["aux_hidden_states"].transpose(0, 1)[: args.seq_length] + batch["hidden_states"] = feature_b["hidden_states"].transpose(0, 1)[: args.seq_length] return batch diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index 1c05aab22dd..7e94a0be111 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -6,6 +6,7 @@ import gc import inspect import json +import math import os import random import sys @@ -16,11 +17,13 @@ from tqdm import tqdm # NOTE: Needs to be before modelopt imports in case megatron.core is not installed. -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) +_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(os.path.abspath(os.path.join(_SCRIPT_DIR, "../../../"))) import modelopt.torch.quantization as mtq from modelopt.recipe import ModelOptPTQRecipe, load_recipe from modelopt.torch.export import import_mcore_gpt_from_hf +from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg from modelopt.torch.utils.dataset_utils import get_dataset_dataloader from modelopt.torch.utils.plugins import megatron_generate, megatron_prefill @@ -29,6 +32,7 @@ # releases. try: from modelopt.torch.utils.plugins.megatron_calibration import ( + get_megatron_calibration_dataloader, get_megatron_calibration_forward_loop, ) @@ -49,8 +53,9 @@ mtq_luts = None warnings.warn("luts is not installed. LUTs quantization configs will not be available.") -from utils import get_hf_tokenizer +from utils import build_lm_batch_from_input_ids, get_hf_tokenizer +from megatron.core import parallel_state from megatron.core.parallel_state import get_context_parallel_group from megatron.core.utils import get_batch_on_this_cp_rank, unwrap_model from megatron.post_training.arguments import add_modelopt_args @@ -65,6 +70,7 @@ warnings.filterwarnings("ignore") + QUANT_CFG_CHOICES = {} # Auto-load all quant configs by full name @@ -128,6 +134,12 @@ def add_text_generate_ptq_args(parser): default=False, help="Skip the post-quantization generate/validation step.", ) + group.add_argument( + "--generate-output-len", + type=int, + default=32, + help="Number of tokens to generate in the post-quantization validation step.", + ) group.add_argument( "--references", type=str, @@ -161,6 +173,45 @@ def add_text_generate_ptq_args(parser): action="store_true", help="Synchronize expert weight amax across experts.", ) + group.add_argument( + "--auto-quantize-bits", + type=float, + default=None, + help=( + "Target effective bits for mtq.auto_quantize per-layer mixed-precision search " + "(e.g. 4.0, 4.8). When set, runs auto-quantize instead of plain mtq.quantize, " + "and --export-quant-cfg / --recipe are ignored." + ), + ) + group.add_argument( + "--auto-quantize-formats", + type=str, + nargs="+", + default=["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"], + help="Quantization format names (entries in mtq.config.choices) to search over.", + ) + group.add_argument( + "--auto-quantize-method", + type=str, + default="gradient", + choices=["gradient", "kl_div"], + help="Method for auto_quantize sensitivity scoring.", + ) + group.add_argument( + "--auto-quantize-score-size", + type=int, + default=128, + help="Number of samples to use for sensitivity scoring in auto_quantize.", + ) + group.add_argument( + "--auto-quantize-checkpoint", + type=str, + default=None, + help=( + "Optional path to save/restore the auto_quantize search state " + "(sensitivity scores, costs, calibration state) across runs." + ), + ) add_modelopt_args(parser) return parser @@ -176,6 +227,48 @@ def check_arguments(): print_rank_0("WARNING: Forcing moe_grouped_gemm to False for PTQ and export.") args.moe_grouped_gemm = False + uses_calibration = args.auto_quantize_bits is not None or ( + (args.export_quant_cfg is not None or args.recipe is not None) and not args.weight_only + ) + if ( + uses_calibration + and args.context_parallel_size > 1 + and args.calib_max_sequence_length % (2 * args.context_parallel_size) != 0 + ): + raise ValueError( + "--calib-max-sequence-length must be a multiple of 2 * " + "--context-parallel-size when context parallelism is enabled." + ) + + if args.auto_quantize_bits is not None and not _HAS_SHARED_CALIB: + raise RuntimeError( + "auto_quantize requires modelopt 0.46+. " + "Upgrade with: pip install nvidia-modelopt>=0.46" + ) + + if args.auto_quantize_bits is not None: + if args.export_quant_cfg is not None: + print_rank_0( + "WARNING: --auto-quantize-bits overrides --export-quant-cfg; the latter is ignored." + ) + args.export_quant_cfg = None + if args.recipe is not None: + print_rank_0( + "WARNING: --auto-quantize-bits overrides --recipe; the latter is ignored." + ) + args.recipe = None + if args.pipeline_model_parallel_size > 1: + raise ValueError( + "auto_quantize currently requires pipeline-model-parallel-size=1 because " + "ModelOpt needs additional support for pipeline parallelism." + ) + for fmt in args.auto_quantize_formats: + if fmt not in QUANT_CFG_CHOICES: + raise ValueError( + f"Unknown auto-quantize format '{fmt}'. Available: " + f"{sorted(QUANT_CFG_CHOICES.keys())}" + ) + def get_modelopt_torch_quantization_config(): """Return a quantization config.""" @@ -312,6 +405,85 @@ def get_calib_dataloader( ) +def auto_quantize_model(unwrapped_model, tokenizer): + """Run mtq.auto_quantize on the MCore model to search per-layer mixed precision. + + Returns the search_state dict produced by mtq.auto_quantize. + """ + args = get_args() + + calib_dataloader = get_megatron_calibration_dataloader( + tokenizer, + dataset_name=args.calib_dataset_path_or_name, + num_samples=args.calib_size, + seq_length=args.calib_max_sequence_length, + batch_size=args.calib_batch_size, + ) + + def forward_step(model, batch): + return megatron_prefill(model, batch["input_ids"]) + + def forward_backward_step(model, batch): + lm_batch = build_lm_batch_from_input_ids( + batch, + cp_group=get_context_parallel_group(), + ) + loss = model.forward( + input_ids=lm_batch["tokens"], + position_ids=lm_batch["position_ids"], + attention_mask=lm_batch["attention_mask"], + labels=lm_batch["labels"], + loss_mask=lm_batch["loss_mask"], + runtime_gather_output=True, + ) + loss.mean().backward() + + quantization_formats = [QUANT_CFG_CHOICES[fmt] for fmt in args.auto_quantize_formats] + disabled_layers = [ + entry["quantizer_name"] + for entry in _default_disabled_quantizer_cfg + if "parent_class" not in entry + ] + + dp_world_size = parallel_state.get_data_parallel_world_size() if torch.distributed.is_initialized() else 1 + num_calib_steps = len(calib_dataloader) + score_samples_per_step = max(dp_world_size * args.calib_batch_size, 1) + num_score_steps = min( + len(calib_dataloader), + max(math.ceil(args.auto_quantize_score_size / score_samples_per_step), 1), + ) + + print_rank_0( + f"Running mtq.auto_quantize: bits={args.auto_quantize_bits}, " + f"formats={args.auto_quantize_formats}, method={args.auto_quantize_method}, " + f"num_calib_steps={num_calib_steps}, num_score_steps={num_score_steps}" + ) + + _, search_state = mtq.auto_quantize( + unwrapped_model, + constraints={"effective_bits": args.auto_quantize_bits}, + quantization_formats=quantization_formats, + data_loader=calib_dataloader, + forward_step=forward_step, + loss_func=None, + forward_backward_step=forward_backward_step, + disabled_layers=disabled_layers, + num_calib_steps=num_calib_steps, + num_score_steps=num_score_steps, + verbose=True, + method=args.auto_quantize_method, + checkpoint=args.auto_quantize_checkpoint, + ) + + if args.save is not None and torch.distributed.get_rank() == 0: + os.makedirs(args.save, exist_ok=True) + torch.save( + search_state, + os.path.join(args.save, f"auto_quantize_search_state_rank_{torch.distributed.get_rank()}.pth"), + ) + return search_state + + if __name__ == "__main__": parse_and_validate_args(extra_args_provider=add_text_generate_ptq_args, args_defaults={ "tokenizer_type": "HuggingFaceTokenizer", @@ -360,10 +532,10 @@ def _custom_prompt_forward_loop_func(model): for idx, prompt in tqdm(enumerate(all_prompts), disable=torch.distributed.get_rank()): tokens = tokenizer(prompt, return_tensors="pt") # enable_kv_cache=False to avoid pre-allocating the static KV cache: this is a - # sanity-check generation (32 tokens), and the KV-cache allocation can OOM tight + # sanity-check generation, and the KV-cache allocation can OOM tight # quantization runs on large MoE models. generated_ids = megatron_generate( - model, tokens.input_ids.cuda(), osl=32, enable_kv_cache=False + model, tokens.input_ids.cuda(), osl=args.generate_output_len, enable_kv_cache=False ) generated_texts = tokenizer.batch_decode(generated_ids) print_rank_0("{}".format(generated_texts)) @@ -399,7 +571,16 @@ def _dataset_forward_loop_func(model): unwrapped_model = unwrap_model(model)[0] - if args.export_quant_cfg is not None or args.recipe is not None: + if args.auto_quantize_bits is not None: + print_rank_0("Running auto-quantize search...") + auto_quantize_model(unwrapped_model, tokenizer) + + if args.compress: + mtq.compress(unwrapped_model) + print_rank_0("Weights are now compressed to low-bit!") + + print_distributed_quant_summary(model, "Auto-Quantized Model:") + elif args.export_quant_cfg is not None or args.recipe is not None: print_rank_0("Quantizing the model...") mtq_config = get_modelopt_torch_quantization_config() diff --git a/examples/post_training/modelopt/quantize.sh b/examples/post_training/modelopt/quantize.sh index e96b224f3c1..c36f3f78b78 100755 --- a/examples/post_training/modelopt/quantize.sh +++ b/examples/post_training/modelopt/quantize.sh @@ -21,9 +21,14 @@ if [ -z ${QUANT_CFG} ]; then fi # If the 2nd positional arg looks like a recipe path (contains '/' or ends in -# '.yaml'/'.yml') pass it via --recipe; otherwise treat it as a built-in +# '.yaml'/'.yml') pass it via --recipe; if it is the literal 'auto' sentinel +# skip the QUANT_CFG flags entirely so --auto-quantize-bits supplied through +# MLM_EXTRA_ARGS drives the search; otherwise treat it as a built-in # config name and pass it via --export-quant-cfg. case "${QUANT_CFG}" in + auto|AUTO|auto_quantize) + QUANT_CFG_ARGS=() + ;; */*|*.yaml|*.yml) QUANT_CFG_ARGS=(--recipe "${QUANT_CFG}") ;; diff --git a/examples/post_training/modelopt/utils.py b/examples/post_training/modelopt/utils.py index fd554caa6d8..512b640fa9c 100644 --- a/examples/post_training/modelopt/utils.py +++ b/examples/post_training/modelopt/utils.py @@ -3,10 +3,15 @@ """Shared utilities for modelopt post-training scripts.""" import os import sys +from typing import Any + +import torch sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) +from megatron.core.utils import get_batch_on_this_cp_rank from megatron.training import get_tokenizer +from megatron.training.utils import get_ltor_masks_and_position_ids def get_hf_tokenizer(): @@ -23,3 +28,142 @@ def get_hf_tokenizer(): tokenizer = getattr(tokenizer, attr) break return tokenizer + + +def get_eos_token_id(hf_tokenizer=None): + """Return the eos token id used for loss and position masking. + + Some tokenizers use eos tokens inside chat turns; this maps known chat eos strings + to the token ids used when packing SFT samples. + """ + if hf_tokenizer is None: + hf_tokenizer = get_hf_tokenizer() + + if hf_tokenizer.eos_token == "<|eot_id|>": + return 128001 + if hf_tokenizer.eos_token == "<|eot|>": + return 200001 + if hf_tokenizer.eos_token == "<|im_end|>": + return 151643 + if hf_tokenizer.eos_token == "<|return|>": + return 199999 + + return hf_tokenizer.eos_token_id + + +def build_lm_batch( + input_ids: torch.Tensor, + seq_length: int, + *, + sample_loss_mask: torch.Tensor | None = None, + pad_attention_mask: torch.Tensor | None = None, + eos_token_id: int | None = None, + reset_position_ids: bool = False, + reset_attention_mask: bool = False, + eod_mask_loss: bool = False, + pad_mask_loss: bool = False, + cp_group: torch.distributed.ProcessGroup | None = None, + is_hybrid_cp: bool = False, +) -> dict[str, torch.Tensor]: + """Build causal-LM training tensors from packed or padded ``input_ids``. + + ``input_ids`` must contain ``seq_length + 1`` tokens per row so that ``tokens`` + and next-token ``labels`` both have length ``seq_length``. + + Args: + input_ids: Token ids with an extra trailing token for the label shift. + seq_length: Number of input tokens (excluding the extra label token). + sample_loss_mask: Optional per-token mask aligned with ``input_ids``. When + provided, only positions with a non-zero mask at the label positions + contribute to ``loss_mask`` (SFT answer-only masking). + pad_attention_mask: Optional HuggingFace-style attention mask aligned with + ``input_ids``. When provided, padding positions are zeroed out in + ``loss_mask`` using the label-aligned slice. + eos_token_id: Eos token id for ``get_ltor_masks_and_position_ids``. + reset_position_ids: Passed through to ``get_ltor_masks_and_position_ids``. + reset_attention_mask: Passed through to ``get_ltor_masks_and_position_ids``. + eod_mask_loss: Passed through to ``get_ltor_masks_and_position_ids``. + pad_mask_loss: Passed through to ``get_ltor_masks_and_position_ids``. + cp_group: When set, slice the batch for context parallelism. + is_hybrid_cp: Passed through to ``get_batch_on_this_cp_rank``. + + Returns: + Dict with ``tokens``, ``labels``, ``loss_mask``, ``attention_mask``, and + ``position_ids`` ready for ``GPTModel.forward``. + """ + if eos_token_id is None: + eos_token_id = get_eos_token_id() + + tokens = input_ids[:, :seq_length].contiguous() + labels = input_ids[:, 1 : seq_length + 1].contiguous() + + attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( + tokens, + eos_token_id, + eos_token_id, + reset_position_ids, + reset_attention_mask, + eod_mask_loss, + pad_mask_loss, + ) + + if sample_loss_mask is not None: + answer_only_loss_mask = sample_loss_mask[:, 1 : seq_length + 1].contiguous() + loss_mask = loss_mask * answer_only_loss_mask.to(dtype=loss_mask.dtype) + + if pad_attention_mask is not None: + pad_mask = pad_attention_mask[:, 1 : seq_length + 1].to(dtype=loss_mask.dtype) + loss_mask = loss_mask * pad_mask + + batch = { + "tokens": tokens, + "labels": labels.contiguous(), + "loss_mask": loss_mask.contiguous(), + "attention_mask": attention_mask, + "position_ids": position_ids, + } + + if cp_group is not None: + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=cp_group) + + return batch + + +def build_lm_batch_from_input_ids( + batch: dict[str, Any], + *, + seq_length: int | None = None, + eos_token_id: int | None = None, + reset_position_ids: bool = False, + reset_attention_mask: bool = False, + eod_mask_loss: bool = False, + pad_mask_loss: bool = False, + cp_group: torch.distributed.ProcessGroup | None = None, + is_hybrid_cp: bool = False, +) -> dict[str, torch.Tensor]: + """Build an LM batch dict from a dataloader batch containing ``input_ids``. + + Calibration and HF dataloaders provide ``input_ids`` of shape + ``[batch, seq_length + 1]`` (or pass ``seq_length=input_ids.shape[1] - 1``). + An optional ``attention_mask`` entry is used to mask padded label positions. + """ + input_ids = batch["input_ids"] + if seq_length is None: + seq_length = input_ids.shape[1] - 1 + + pad_attention_mask = batch.get("attention_mask") + sample_loss_mask = batch.get("loss_mask") + + return build_lm_batch( + input_ids, + seq_length, + sample_loss_mask=sample_loss_mask, + pad_attention_mask=pad_attention_mask, + eos_token_id=eos_token_id, + reset_position_ids=reset_position_ids, + reset_attention_mask=reset_attention_mask, + eod_mask_loss=eod_mask_loss, + pad_mask_loss=pad_mask_loss, + cp_group=cp_group, + is_hybrid_cp=is_hybrid_cp, + )