diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index cc1de2a0b..b595de990 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -5,7 +5,7 @@ import torch from torch._dynamo import config as dynamo_config import torch.distributed as dist -from torch.nn.attention.flex_attention import AuxOutput, AuxRequest, BlockMask +from torch.nn.attention.flex_attention import BlockMask import triton import triton.language as tl @@ -16,6 +16,7 @@ get_sparse_compiled_flex_attention, normalize_flex_lse, normalize_sparse_block_size, + prepare_sparse_flex_attention, select_sparse_execution_family, sparse_compiled_flex_attention, ) @@ -23,6 +24,7 @@ from .block_mask import build_block_mask_from_context, prepare_block_mask_context from .comm import A2AVCommunicator from .range_ops import ( + prepare_range_head_major_kernels, range_gather_head_major, range_reduce_sum_, range_reduce_sum_head_major_, @@ -688,21 +690,21 @@ def run( triton_num_stages_2_head_dims=self.triton_num_stages_2_head_dims, ) ) - out, aux = cast( - tuple[torch.Tensor, AuxOutput], - compiled_flex_attention( - q, - k, - v, - block_mask=block_mask, - scale=scale, - enable_gqa=enable_gqa, - return_aux=AuxRequest(lse=True), - ), + prepare_sparse_flex_attention( + q, + k, + v, + block_mask=block_mask, + enable_gqa=enable_gqa, + ) + out, lse = compiled_flex_attention( + q, + k, + v, + block_mask=block_mask, + scale=scale, + enable_gqa=enable_gqa, ) - lse = aux.lse - if lse is None: - raise RuntimeError("Compiled flex attention did not return lse.") lse = normalize_flex_lse(lse, backend=backend) return out, lse @@ -1021,9 +1023,9 @@ def _run_stage_attention( head_major=input_head_major, ) if input_head_major: - q_flex = q_stage.unsqueeze(0) - k_flex = k_stage.unsqueeze(0) - v_flex = v_stage.unsqueeze(0) + q_flex = q_stage.unsqueeze(0).contiguous() + k_flex = k_stage.unsqueeze(0).contiguous() + v_flex = v_stage.unsqueeze(0).contiguous() else: q_flex = q_stage.permute(1, 0, 2).unsqueeze(0).contiguous() k_flex = k_stage.permute(1, 0, 2).unsqueeze(0).contiguous() @@ -1887,11 +1889,14 @@ def _flatten_qkv( value: torch.Tensor, state: ArtContextParallelState, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return ( + flattened = ( flatten_valid_sequence_head_major(query, state.rank_plan.local_valid_lengths), flatten_valid_sequence_head_major(key, state.rank_plan.local_valid_lengths), flatten_valid_sequence_head_major(value, state.rank_plan.local_valid_lengths), ) + for tensor in flattened: + prepare_range_head_major_kernels(tensor) + return flattened def _run_context_parallel_forward( diff --git a/src/art/megatron/context_parallel/range_ops.py b/src/art/megatron/context_parallel/range_ops.py index a645e9f80..7a0ef8b6d 100644 --- a/src/art/megatron/context_parallel/range_ops.py +++ b/src/art/megatron/context_parallel/range_ops.py @@ -8,6 +8,8 @@ from .types import TokenRange +_WARMED_HEAD_MAJOR_KERNELS: set[tuple[str, int | None, torch.dtype, int]] = set() + def _single_range(ranges: Sequence[TokenRange]) -> TokenRange | None: compact = [range_ for range_ in ranges if range_.size() > 0] @@ -90,8 +92,7 @@ def _range_gather_head_major_kernel( output_head_stride, output_token_stride, inner_size: tl.constexpr, - n_cols: tl.constexpr, - n_col_blocks: tl.constexpr, + n_cols, elem_per_block: tl.constexpr, ): out_row = tl.program_id(0) @@ -129,8 +130,7 @@ def _range_reduce_sum_head_major_kernel( output_head_stride, output_token_stride, inner_size: tl.constexpr, - n_cols: tl.constexpr, - n_col_blocks: tl.constexpr, + n_cols, elem_per_block: tl.constexpr, ): in_row = tl.program_id(0) @@ -196,7 +196,7 @@ def _range_meta( dtype=torch.int64, ) cu_range_sizes = torch.cumsum(range_sizes, dim=0) - total_size = int(cu_range_sizes[-1].item()) + total_size = sum(range_.size() for range_ in compact) row_map = torch.repeat_interleave( torch.arange(len(compact), device=device, dtype=torch.int64), range_sizes[1:], @@ -382,14 +382,46 @@ def _range_gather_head_major_impl( output.stride(0), output.stride(1), inner_size=inner_size, # ty: ignore[invalid-argument-type] - n_cols=n_cols, # ty: ignore[invalid-argument-type] - n_col_blocks=n_col_blocks, + n_cols=n_cols, elem_per_block=elem_per_block, # ty: ignore[invalid-argument-type] num_warps=4, # ty: ignore[unknown-argument] ) return output +def prepare_range_head_major_kernels(input_tensor: torch.Tensor) -> None: + """Compile one shape-polymorphic gather/reduce pair before variable CP plans. + + Prefix trees may not require multi-range Q/KV movement until a late batch. + Warm once per device, dtype, and inner dimension so that workload variation + cannot expose compilation in a later optimizer step. This enqueues two tiny + kernels on the current stream and must remain free of CUDA synchronization. + """ + if not input_tensor.is_cuda: + return + if input_tensor.ndim < 3: + raise RuntimeError( + "Head-major range kernel preparation expects [H, T, ...], " + f"got {tuple(input_tensor.shape)}" + ) + inner_size = input_tensor.numel() // max( + int(input_tensor.shape[0] * input_tensor.shape[1]), 1 + ) + key = ( + input_tensor.device.type, + input_tensor.device.index, + input_tensor.dtype, + inner_size, + ) + if key in _WARMED_HEAD_MAJOR_KERNELS: + return + probe = input_tensor.new_empty((input_tensor.shape[0], 2, *input_tensor.shape[2:])) + ranges = (TokenRange(start=0, end=1), TokenRange(start=1, end=2)) + gathered = _range_gather_head_major_impl(probe, ranges) + range_reduce_sum_head_major_(gathered, output_tensor=probe, ranges=ranges) + _WARMED_HEAD_MAJOR_KERNELS.add(key) + + class _RangeGatherFn(torch.autograd.Function): @staticmethod def forward( @@ -675,8 +707,7 @@ def range_reduce_sum_head_major_( output_tensor.stride(0), output_tensor.stride(1), inner_size=inner_size, # ty: ignore[invalid-argument-type] - n_cols=n_cols, # ty: ignore[invalid-argument-type] - n_col_blocks=n_col_blocks, + n_cols=n_cols, elem_per_block=elem_per_block, # ty: ignore[invalid-argument-type] num_warps=4, # ty: ignore[unknown-argument] ) diff --git a/src/art/megatron/flex_attn/compiled.py b/src/art/megatron/flex_attn/compiled.py index b6bd8dfe2..00b220d96 100644 --- a/src/art/megatron/flex_attn/compiled.py +++ b/src/art/megatron/flex_attn/compiled.py @@ -4,8 +4,12 @@ from typing import Any, Literal, TypeAlias, cast import torch +from torch._higher_order_ops.flex_attention import ( + flex_attention as flex_attention_hop, +) from torch.nn.attention.flex_attention import ( AuxRequest, + BlockMask, FlexKernelOptions, flex_attention, ) @@ -105,26 +109,72 @@ def _forced_flex_attention_dense( ) -def _forced_flex_attention_sparse( - q, - k, - v, +def _identity_score(score, _batch, _head, _query, _key): + return score + + +def _sparse_flex_attention_with_options(kernel_options: FlexKernelOptions) -> Any: + resolved_options = dict(kernel_options) + resolved_options.setdefault("PRESCALE_QK", False) + resolved_options.setdefault("ROWS_GUARANTEED_SAFE", False) + resolved_options.setdefault("BLOCKS_ARE_CONTIGUOUS", False) + resolved_options.setdefault("WRITE_DQ", True) + resolved_options["OUTPUT_LOGSUMEXP"] = True + resolved_options["OUTPUT_MAX"] = False + + def _flex_attention(q, k, v, *, block_mask, scale, enable_gqa): + del enable_gqa + for tensor in (q, k, v): + torch._dynamo.mark_static(tensor, -3) + torch._dynamo.mark_static(tensor, -1) + out, lse, _max_scores = flex_attention_hop( + q, + k, + v, + _identity_score, + (q.shape[-2], k.shape[-2], *block_mask.as_tuple()[2:]), + scale, + resolved_options, + ) + return out, lse * _FLASH_LSE_RESCALE + + return _flex_attention + + +def prepare_sparse_flex_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, *, - block_mask, - scale, - enable_gqa, - return_aux: AuxRequest | None = None, -): - return flex_attention( - q, - k, - v, - block_mask=block_mask, - scale=scale, - enable_gqa=enable_gqa, - kernel_options=_FORCED_FLEX_KERNEL_OPTIONS, - return_aux=return_aux, - ) + block_mask: BlockMask, + enable_gqa: bool, +) -> None: + """Validate ART's sparse boundary and keep stage lengths symbolic.""" + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise RuntimeError("Sparse flex attention requires rank-4 q, k, and v.") + if int(q.shape[0]) != int(k.shape[0]) or k.shape[:-1] != v.shape[:-1]: + raise RuntimeError("Sparse flex attention received incompatible q, k, and v.") + if int(q.shape[-1]) != int(k.shape[-1]): + raise RuntimeError("Sparse flex attention requires equal q/k head dimensions.") + q_heads, kv_heads = int(q.shape[1]), int(k.shape[1]) + if (q_heads != kv_heads and not enable_gqa) or q_heads % kv_heads != 0: + raise RuntimeError( + f"Sparse flex attention received q_heads={q_heads}, kv_heads={kv_heads}, " + f"enable_gqa={enable_gqa}." + ) + if tuple(block_mask.shape[-2:]) != (int(q.shape[2]), int(k.shape[2])): + raise RuntimeError("Sparse flex attention block-mask lengths do not match q/k.") + for tensor in (q, k, v): + torch._dynamo.mark_dynamic(tensor, 2) + for value in block_mask.as_tuple()[2:]: + if isinstance(value, torch.Tensor): + for dim in range(2, value.ndim): + torch._dynamo.mark_dynamic(value, dim) + for cell in getattr(block_mask.mask_mod, "__closure__", None) or (): + value = cell.cell_contents + if isinstance(value, torch.Tensor): + for dim in range(value.ndim): + torch._dynamo.mark_dynamic(value, dim) def _flex_attention_with_options(kernel_options: FlexKernelOptions) -> Any: @@ -161,11 +211,16 @@ def select_sparse_execution_family( ) -> tuple[int, int, str]: del is_local_stage q_block, k_block = normalize_sparse_block_size(block_size) + # Avoid Flex's separate zero/one-block Dynamo specialization. target_q_len = ( - 0 if int(q_len) <= 0 else ((int(q_len) + q_block - 1) // q_block) * q_block + 0 + if int(q_len) <= 0 + else max(2, (int(q_len) + q_block - 1) // q_block) * q_block ) target_k_len = ( - 0 if int(k_len) <= 0 else ((int(k_len) + k_block - 1) // k_block) * k_block + 0 + if int(k_len) <= 0 + else max(2, (int(k_len) + k_block - 1) // k_block) * k_block ) return int(target_q_len), int(target_k_len), "sparse" @@ -243,14 +298,14 @@ def get_sparse_compiled_flex_attention( ) sparse_compiled_flex_attention = torch.compile( - _forced_flex_attention_sparse, + _sparse_flex_attention_with_options(_FORCED_FLEX_KERNEL_OPTIONS), ) flash_sparse_compiled_flex_attention = torch.compile( - _flex_attention_with_options(_FLASH_FLEX_KERNEL_OPTIONS), + _sparse_flex_attention_with_options(_FLASH_FLEX_KERNEL_OPTIONS), ) triton_sparse_compiled_flex_attention = torch.compile( - _flex_attention_with_options(_TRITON_FLEX_KERNEL_OPTIONS), + _sparse_flex_attention_with_options(_TRITON_FLEX_KERNEL_OPTIONS), ) triton_num_stages_2_sparse_compiled_flex_attention = torch.compile( - _flex_attention_with_options(_TRITON_NUM_STAGES_2_FLEX_KERNEL_OPTIONS), + _sparse_flex_attention_with_options(_TRITON_NUM_STAGES_2_FLEX_KERNEL_OPTIONS), ) diff --git a/src/art/megatron/model_support/__init__.py b/src/art/megatron/model_support/__init__.py index 7e1a8b374..6a5735363 100644 --- a/src/art/megatron/model_support/__init__.py +++ b/src/art/megatron/model_support/__init__.py @@ -4,6 +4,8 @@ GEMMA4_DENSE_SPEC, GEMMA4_MOE_MODELS, GEMMA4_MOE_SPEC, + LLAMA3_DENSE_MODELS, + LLAMA3_DENSE_SPEC, PROBE_ONLY_MODEL_SUPPORT_SPECS, QWEN3_5_DENSE_MODELS, QWEN3_5_DENSE_SPEC, @@ -66,6 +68,8 @@ def __getattr__(name: str): "GEMMA4_MOE_MODELS", "GEMMA4_MOE_SPEC", "LayerFamilyInstance", + "LLAMA3_DENSE_MODELS", + "LLAMA3_DENSE_SPEC", "ModelSupportHandler", "ModelSupportSpec", "NativeVllmLoraStatus", diff --git a/src/art/megatron/model_support/handlers/__init__.py b/src/art/megatron/model_support/handlers/__init__.py index a8da0158d..670468a12 100644 --- a/src/art/megatron/model_support/handlers/__init__.py +++ b/src/art/megatron/model_support/handlers/__init__.py @@ -16,6 +16,14 @@ "art.megatron.model_support.handlers.default_dense", "DefaultMoeHandler", ), + "LLAMA3_DENSE_HANDLER": ( + "art.megatron.model_support.handlers.llama3", + "LLAMA3_DENSE_HANDLER", + ), + "Llama3DenseHandler": ( + "art.megatron.model_support.handlers.llama3", + "Llama3DenseHandler", + ), "QWEN3_DENSE_HANDLER": ( "art.megatron.model_support.handlers.qwen3_dense", "QWEN3_DENSE_HANDLER", diff --git a/src/art/megatron/model_support/handlers/llama3.py b/src/art/megatron/model_support/handlers/llama3.py new file mode 100644 index 000000000..6ab3783fc --- /dev/null +++ b/src/art/megatron/model_support/handlers/llama3.py @@ -0,0 +1,21 @@ +from typing import Any, Sequence + +from art.megatron.model_support.handlers.default_dense import DefaultDenseHandler +from art.megatron.model_support.handlers.qwen3_common import ( + install_qwen3_text_preprocess_patch, + qwen3_forward_kwargs, +) + + +class Llama3DenseHandler(DefaultDenseHandler): + key = "llama3_dense" + native_vllm_lora_status = "validated" + + def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None: + install_qwen3_text_preprocess_patch(model_chunks) + + def get_forward_kwargs(self, model: Any, **kwargs: Any) -> dict[str, Any]: + return qwen3_forward_kwargs(model, **kwargs) + + +LLAMA3_DENSE_HANDLER = Llama3DenseHandler() diff --git a/src/art/megatron/model_support/registry.py b/src/art/megatron/model_support/registry.py index 0c9926ab6..5511205a9 100644 --- a/src/art/megatron/model_support/registry.py +++ b/src/art/megatron/model_support/registry.py @@ -8,6 +8,7 @@ ) _DEFAULT_DENSE_HANDLER_KEY = "default_dense" +_LLAMA3_DENSE_HANDLER_KEY = "llama3_dense" _QWEN3_DENSE_HANDLER_KEY = "qwen3_dense" _QWEN3_MOE_HANDLER_KEY = "qwen3_moe" _QWEN3_5_DENSE_HANDLER_KEY = "qwen3_5_dense" @@ -81,6 +82,30 @@ native_vllm_lora_status=_DISABLED_NATIVE_VLLM_LORA_STATUS, ) +LLAMA3_DENSE_SPEC = ModelSupportSpec( + key="llama3_dense", + handler_key=_LLAMA3_DENSE_HANDLER_KEY, + model_names=( + "meta-llama/Meta-Llama-3-8B", + "meta-llama/Meta-Llama-3-8B-Instruct", + "meta-llama/Meta-Llama-3-70B", + "meta-llama/Meta-Llama-3-70B-Instruct", + "meta-llama/Llama-3.1-8B", + "meta-llama/Llama-3.1-8B-Instruct", + "meta-llama/Llama-3.1-70B", + "meta-llama/Llama-3.1-70B-Instruct", + "meta-llama/Llama-3.1-405B", + "meta-llama/Llama-3.1-405B-Instruct", + "meta-llama/Llama-3.2-1B", + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B", + "meta-llama/Llama-3.2-3B-Instruct", + "meta-llama/Llama-3.3-70B-Instruct", + ), + default_target_modules=_DENSE_TARGET_MODULES, + native_vllm_lora_status=_VALIDATED_NATIVE_VLLM_LORA_STATUS, +) + QWEN3_MOE_SPEC = ModelSupportSpec( key="qwen3_moe", handler_key=_QWEN3_MOE_HANDLER_KEY, @@ -212,6 +237,7 @@ ) VALIDATED_MODEL_SUPPORT_SPECS = ( + LLAMA3_DENSE_SPEC, QWEN3_MOE_SPEC, QWEN3_DENSE_SPEC, QWEN3_5_MOE_SPEC, @@ -243,6 +269,10 @@ "art.megatron.model_support.handlers.default_dense", "DEFAULT_DENSE_HANDLER", ), + _LLAMA3_DENSE_HANDLER_KEY: ( + "art.megatron.model_support.handlers.llama3", + "LLAMA3_DENSE_HANDLER", + ), _QWEN3_DENSE_HANDLER_KEY: ( "art.megatron.model_support.handlers.qwen3_dense", "QWEN3_DENSE_HANDLER", @@ -302,6 +332,7 @@ _REGISTERED_BRIDGE_KEYS: set[str] = set() QWEN3_DENSE_MODELS = frozenset(QWEN3_DENSE_SPEC.model_names) +LLAMA3_DENSE_MODELS = frozenset(LLAMA3_DENSE_SPEC.model_names) QWEN3_MOE_MODELS = frozenset(QWEN3_MOE_SPEC.model_names) QWEN3_5_DENSE_MODELS = frozenset(QWEN3_5_DENSE_SPEC.model_names) QWEN3_5_MOE_MODELS = frozenset(QWEN3_5_MOE_SPEC.model_names) diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index fe04050f9..92f93c2aa 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -199,9 +199,6 @@ def _skip_meta_to( adapter_config=final_config, ) del peft_model, model - if torch.cuda.is_available(): - torch.cuda.synchronize() - torch.cuda.empty_cache() @dataclass diff --git a/src/art/megatron/weights/adapter_export.py b/src/art/megatron/weights/adapter_export.py index 664e30a02..bb96b51c7 100644 --- a/src/art/megatron/weights/adapter_export.py +++ b/src/art/megatron/weights/adapter_export.py @@ -15,6 +15,7 @@ SelfAttentionLinearProjLoRA, SelfAttentionLinearQKVLoRA, SharedExpertsLinearFC1LoRA, + SharedExpertsLinearFC2LoRA, ) from art.megatron.weights.param_name_canonicalization import canonical_art_param_name @@ -425,10 +426,10 @@ def add_split_mlp_adapter_weights( ) linear_fc2 = getattr(mlp, "linear_fc2", None) - if isinstance(linear_fc2, SelfAttentionLinearProjLoRA): + if isinstance(linear_fc2, SharedExpertsLinearFC2LoRA): fc2_prefix = f"{base_prefix}.linear_fc2" _set_adapter_weights( adapter_weights_by_base, fc2_prefix, - _simple_adapter_weight(fc2_prefix, linear_fc2.lora), + _simple_adapter_weight(fc2_prefix, linear_fc2.row_parallel_lora.lora), ) diff --git a/tests/integration/megatron/model_support/oracle_worker.py b/tests/integration/megatron/model_support/oracle_worker.py index 7904c1af5..014871667 100644 --- a/tests/integration/megatron/model_support/oracle_worker.py +++ b/tests/integration/megatron/model_support/oracle_worker.py @@ -505,6 +505,24 @@ def _configure_cuda_precision(case_config: OracleCaseConfig) -> None: torch.set_float32_matmul_precision("highest") +def _sparse_flex_test_call(attention_call: Callable[..., Any]) -> Callable[..., Any]: + from torch.nn.attention.flex_attention import AuxRequest + + def sparse_call(q, k, v, *, block_mask, scale, enable_gqa): + out, aux = attention_call( + q, + k, + v, + block_mask=block_mask, + scale=scale, + enable_gqa=enable_gqa, + return_aux=AuxRequest(lse=True), + ) + return out, aux.lse + + return sparse_call + + @contextmanager def _apply_requested_flex_backend_patch(flex_backend: str | None): if flex_backend is None: @@ -539,7 +557,9 @@ def _apply_requested_flex_backend_patch(flex_backend: str | None): compiled_flex_attention._forced_flex_attention_dense ) compiled_flex_attention.sparse_compiled_flex_attention = torch.compile( - compiled_flex_attention._forced_flex_attention_sparse + compiled_flex_attention._sparse_flex_attention_with_options( + patched_kernel_options + ) ) try: yield @@ -594,7 +614,7 @@ def _fp32_inner_call( _fp32_inner_call ) compiled_flex_attention.sparse_compiled_flex_attention = torch.compile( - _fp32_inner_call + _sparse_flex_test_call(_fp32_inner_call) ) try: yield @@ -678,7 +698,7 @@ def _attention_forward_fp32(self, hidden_states, *args, **kwargs): _fp32_inner_call ) compiled_flex_attention.sparse_compiled_flex_attention = torch.compile( - _fp32_inner_call + _sparse_flex_test_call(_fp32_inner_call) ) setattr(ColumnParallelLinear, "_forward_impl", _column_forward_impl_fp32) setattr(RowParallelLinear, "_forward_impl", _row_forward_impl_fp32) diff --git a/tests/integration/megatron/model_support/packing_invariance.py b/tests/integration/megatron/model_support/packing_invariance.py index ca72a2684..68be89ada 100644 --- a/tests/integration/megatron/model_support/packing_invariance.py +++ b/tests/integration/megatron/model_support/packing_invariance.py @@ -49,19 +49,6 @@ PACKING_INVARIANCE_REPORT_FILENAME = "report.json" PACKING_INVARIANCE_ARTIFACT_SUITE_NAME = "Megatron packing-invariance artifacts" REPO_ROOT = Path(__file__).resolve().parents[4] -_SINGLE_ROTARY_OUTPUT_HANDLER_KEYS = frozenset( - { - "default_dense", - "default_moe", - "qwen3_dense", - "qwen3_moe", - "qwen3_5_dense", - "qwen3_5_moe", - "dsv4", - "gpt_oss_moe", - } -) -_TUPLE_ROTARY_OUTPUT_HANDLER_KEYS = frozenset({"gemma4_dense", "gemma4_moe"}) def _slugify(value: str) -> str: @@ -283,21 +270,18 @@ def _rotary_grouping_check( def _rotary_outputs_for_validation( *, - handler_key: str, preprocess_output: Any, ) -> tuple[torch.Tensor | None, ...]: rotary_output = preprocess_output[1] - if handler_key in _SINGLE_ROTARY_OUTPUT_HANDLER_KEYS: - return ( - cast(torch.Tensor | None, rotary_output) - if torch.is_tensor(rotary_output) - else None, - ) - if handler_key in _TUPLE_ROTARY_OUTPUT_HANDLER_KEYS: - local_rotary, global_rotary = rotary_output - return local_rotary, global_rotary + if rotary_output is None or torch.is_tensor(rotary_output): + return (cast(torch.Tensor | None, rotary_output),) + if isinstance(rotary_output, tuple) and all( + item is None or torch.is_tensor(item) for item in rotary_output + ): + return cast(tuple[torch.Tensor | None, ...], rotary_output) raise RuntimeError( - f"Packed position validation has no rotary output mapping for {handler_key!r}" + "Packed position validation received unsupported rotary outputs: " + f"{type(rotary_output).__name__}" ) @@ -703,7 +687,6 @@ def _run_packing_invariance_worker( row_respected = True row_repeated_count = 0 rotary_outputs = _rotary_outputs_for_validation( - handler_key=runtime.model_support_handler.key, preprocess_output=hooked_output, ) for rotary_output in rotary_outputs: diff --git a/tests/integration/megatron/model_support/test_provider_support.py b/tests/integration/megatron/model_support/test_provider_support.py index 708d9f6c9..43a641452 100644 --- a/tests/integration/megatron/model_support/test_provider_support.py +++ b/tests/integration/megatron/model_support/test_provider_support.py @@ -136,7 +136,19 @@ def test_openpipe_qwen3_14b_instruct_uses_qwen3_dense_support() -> None: assert handler.key == "qwen3_dense" +def test_meta_llama_32_1b_instruct_uses_llama3_dense_support() -> None: + model = "meta-llama/Llama-3.2-1B-Instruct" + spec = get_model_support_spec(model) + handler = get_model_support_handler(model) + + assert spec.key == "llama3_dense" + assert spec.is_moe is False + assert spec.native_vllm_lora_status == "validated" + assert handler.key == "llama3_dense" + + def test_model_support_specs_own_moe_metadata() -> None: + assert model_uses_expert_parallel("meta-llama/Llama-3.2-1B-Instruct") is False assert model_uses_expert_parallel("OpenPipe/Qwen3-14B-Instruct") is False assert model_uses_expert_parallel("Qwen/Qwen3-30B-A3B-Instruct-2507") is True assert model_uses_expert_parallel("Qwen/Qwen3.5-35B-A3B") is True diff --git a/tests/integration/megatron/model_support/test_workflow.py b/tests/integration/megatron/model_support/test_workflow.py index 84b796937..4b68579cd 100644 --- a/tests/integration/megatron/model_support/test_workflow.py +++ b/tests/integration/megatron/model_support/test_workflow.py @@ -73,6 +73,7 @@ def test_build_validation_stage_names_has_fixed_order() -> None: def test_validated_architecture_representative_models_are_fixed() -> None: assert validated_architecture_representative_models() == [ + "meta-llama/Llama-3.2-1B-Instruct", "Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-32B", "Qwen/Qwen3.5-35B-A3B", @@ -243,7 +244,11 @@ def _build_validation_report( stop_on_failure=True, ) - assert calls == ["Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-32B"] + assert calls == [ + "meta-llama/Llama-3.2-1B-Instruct", + "Qwen/Qwen3-30B-A3B", + "Qwen/Qwen3-32B", + ] assert report.passed is False assert [item.base_model for item in report.reports] == calls diff --git a/tests/integration/megatron/model_support/workflow.py b/tests/integration/megatron/model_support/workflow.py index f42d3ed21..3a4350930 100644 --- a/tests/integration/megatron/model_support/workflow.py +++ b/tests/integration/megatron/model_support/workflow.py @@ -62,6 +62,7 @@ ) ALL_VALIDATION_STAGES = (*MANDATORY_VALIDATION_STAGES, *OPTIONAL_VALIDATION_STAGES) ARCHITECTURE_REPRESENTATIVE_MODELS = { + "llama3_dense": "meta-llama/Llama-3.2-1B-Instruct", "qwen3_moe": "Qwen/Qwen3-30B-A3B", "qwen3_dense": "Qwen/Qwen3-32B", "qwen3_5_moe": "Qwen/Qwen3.5-35B-A3B", diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py index 11df05a1a..b97e92fc4 100644 --- a/tests/integration/megatron/train_inf_mismatch/real_path.py +++ b/tests/integration/megatron/train_inf_mismatch/real_path.py @@ -401,29 +401,28 @@ async def _rollout( import art messages = [{"role": "user", "content": prompt}] - - async def _request() -> None: - request_kwargs: dict[str, Any] = {} - if extra_body is not None: - request_kwargs["extra_body"] = extra_body - response = await model.openai_client().chat.completions.create( - model=model.get_inference_name(), - messages=messages, - max_tokens=max_completion_tokens, - temperature=0.8, - logprobs=True, - top_logprobs=TOP_K, - **request_kwargs, - ) - if trajectory := art.auto_trajectory(): # ty: ignore[deprecated] - logprobs = response.choices[0].logprobs - trajectory.reward = reward - trajectory.metrics["completion_tokens"] = ( + request_kwargs: dict[str, Any] = {} + if extra_body is not None: + request_kwargs["extra_body"] = extra_body + response = await model.openai_client().chat.completions.create( + model=model.get_inference_name(), + messages=messages, + max_tokens=max_completion_tokens, + temperature=0.8, + logprobs=True, + top_logprobs=TOP_K, + **request_kwargs, + ) + choice = response.choices[0] + logprobs = choice.logprobs + return art.Trajectory( + messages_and_choices=[*messages, choice], + reward=reward, + metrics={ + "completion_tokens": ( len(logprobs.content or []) if logprobs is not None else 0 ) - - return await art.capture_auto_trajectory( # ty: ignore[deprecated] - _request() + }, ) diff --git a/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py b/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py index 65c2b449e..1e68277fb 100644 --- a/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py +++ b/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py @@ -1,7 +1,11 @@ from __future__ import annotations import math +from types import SimpleNamespace +from unittest.mock import AsyncMock +from openai.types.chat.chat_completion import Choice +from openai.types.chat.chat_completion_message import ChatCompletionMessage import pytest torch = pytest.importorskip("torch") @@ -28,6 +32,7 @@ _delete_adapter_safetensors_on_pass, _real_path_rollout_mode, _real_path_rollout_weights_mode, + _rollout, ) @@ -192,6 +197,47 @@ def test_real_path_default_generates_16_tokens_per_rollout() -> None: assert RealPathConfig().max_completion_tokens == 16 +@pytest.mark.asyncio +async def test_real_path_rollout_builds_explicit_legacy_trajectory() -> None: + choice = Choice( + index=0, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="answer"), + ) + create = AsyncMock(return_value=SimpleNamespace(choices=[choice])) + model = SimpleNamespace( + get_inference_name=lambda: "test-model", + openai_client=lambda: SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ), + ) + + trajectory = await _rollout( + model=model, + prompt="question", + max_completion_tokens=7, + reward=0.75, + extra_body={"fixture": True}, + ) + + assert trajectory.messages_and_choices == [ + {"role": "user", "content": "question"}, + choice, + ] + assert trajectory.exchanges.chat_completions == [] + assert trajectory.reward == 0.75 + assert trajectory.metrics["completion_tokens"] == 0 + create.assert_awaited_once_with( + model="test-model", + messages=[{"role": "user", "content": "question"}], + max_tokens=7, + temperature=0.8, + logprobs=True, + top_logprobs=TOP_K, + extra_body={"fixture": True}, + ) + + def test_real_path_rollout_mode_follows_config() -> None: native_config = TrainInfOutputParityConfig( base_model="Qwen/Qwen3.5-35B-A3B", diff --git a/tests/integration/megatron/trainability/test_config.py b/tests/integration/megatron/trainability/test_config.py index 5262bb8e7..c68b7e210 100644 --- a/tests/integration/megatron/trainability/test_config.py +++ b/tests/integration/megatron/trainability/test_config.py @@ -129,7 +129,7 @@ def test_megatron_variants_keep_short_packed_sequence_default(monkeypatch) -> No ) assert _variant_packed_sequence_length(variant) == 1024 - assert _variant_train_kwargs(variant) == {"packed_sequence_length": 1024} + assert _variant_train_kwargs(variant) == {} config = _build_internal_config( variant, base_model="Qwen/Qwen3-30B-A3B-Instruct-2507" ) @@ -217,7 +217,7 @@ def test_yes_no_trainability_passes_initially_saturated_stable_report() -> None: step=1, eval_reward=0.9375, train_reward=0.875, - train_metrics={"grad_norm": 54.0}, + train_metrics={"loss/grad_norm": 54.0}, ) ], ) diff --git a/tests/integration/megatron/trainability/yes_no_trainability.py b/tests/integration/megatron/trainability/yes_no_trainability.py index 7343b813a..a0718bf27 100644 --- a/tests/integration/megatron/trainability/yes_no_trainability.py +++ b/tests/integration/megatron/trainability/yes_no_trainability.py @@ -50,7 +50,7 @@ _RESOURCE_STAGE_NAME = Literal["yes_no_trainability", "length_trainability"] -class _TrainKwargs(TypedDict): +class _TrainKwargs(TypedDict, total=False): packed_sequence_length: int @@ -550,6 +550,8 @@ def _variant_packed_sequence_length(variant: _TrainabilityVariant) -> int: def _variant_train_kwargs(variant: _TrainabilityVariant) -> _TrainKwargs: + if variant.backend_name == "megatron": + return {} return {"packed_sequence_length": _variant_packed_sequence_length(variant)} @@ -1035,7 +1037,7 @@ async def run_yes_no_trainability_async( 1e-4, ), loss_fn="cispo", - packed_sequence_length=train_kwargs["packed_sequence_length"], + **train_kwargs, ) await model.log( train_groups, @@ -1119,7 +1121,9 @@ def yes_no_trainability_passed(report: YesNoTrainabilityReport) -> bool: and report.final_eval_reward is not None and report.final_eval_reward >= report.reward_threshold and bool(report.steps) - and any(step.train_metrics.get("grad_norm", 0.0) > 0.0 for step in report.steps) + and any( + step.train_metrics.get("loss/grad_norm", 0.0) > 0.0 for step in report.steps + ) ) return learned_from_below_threshold or already_saturated_and_stable