feat(qwen): support Qwen3.5 and Qwen3-VL-MoE parallel training-hyper-parallel-AtomGit

Mhyper_parallel/core/pipeline_parallel/scheduler.py+4-1

	@@ -409,7 +409,10 @@ class PipelineScheduleRuntime(ABC):
409	managed_stage_indices = {	409	managed_stage_indices = {
410	stage.stage_index	410	stage.stage_index
411	for stage in self.stages	411	for stage in self.stages
412	- if isinstance(stage.submodule, HSDPModule)	412	+ if any(
		413	+ isinstance(module, HSDPModule)
		414	+ for _, module in platform.get_cells_and_names(stage.submodule)
		415	+ )
413	}	416	}
414	if not managed_stage_indices:	417	if not managed_stage_indices:
415	return	418	return

Mhyper_parallel/core/pipeline_parallel/stage.py+104-26

	@@ -34,8 +34,16 @@ class SharedParameterInfo:
34	Args:	34	Args:
35	parameter (Parameter): The shared parameter object.	35	parameter (Parameter): The shared parameter object.
36	shared_stage (list): The shared stage list.	36	shared_stage (list): The shared stage list.
		37	+ owner_module (Module, optional): Module owning the parameter. When given
		38	+ with ``param_name``, the current parameter is re-fetched from it each
		39	+ access, so a meta-init pipeline that materializes the stage after
		40	+ this info is created (``to_empty`` replaces the Parameter object)
		41	+ still sees the live parameter (and its gradient) rather than a stale
		42	+ orphan. Default ``None`` keeps the parameter object as-is.
		43	+ param_name (str, optional): Attribute name of the parameter on
		44	+ ``owner_module`` (e.g. ``"weight"``).
37	"""	45	"""
38	- def __init__(self, parameter, shared_stage):	46	+ def __init__(self, parameter, shared_stage, owner_module=None, param_name=None):
39	if not isinstance(parameter, platform.Parameter):	47	if not isinstance(parameter, platform.Parameter):
40	raise TypeError(f"Argument 'parameter' must be type of Parameter, \	48	raise TypeError(f"Argument 'parameter' must be type of Parameter, \
41	but got type {type(parameter)}.")	49	but got type {type(parameter)}.")
	@@ -44,11 +52,15 @@ class SharedParameterInfo:
44	but got type {type(shared_stage)}.")	52	but got type {type(shared_stage)}.")
45	self._shared_stage = shared_stage	53	self._shared_stage = shared_stage
46	self._parameter = parameter	54	self._parameter = parameter
		55	+ self._owner_module = owner_module
		56	+ self._param_name = param_name
47	self.group = None	57	self.group = None
48		58
49	@property	59	@property
50	def parameter(self):	60	def parameter(self):
51	"""Return the shared parameter object."""	61	"""Return the shared parameter object."""
		62	+ if self._owner_module is not None and self._param_name is not None:
		63	+ return getattr(self._owner_module, self._param_name)
52	return self._parameter	64	return self._parameter
53		65
54	@property	66	@property
	@@ -178,9 +190,16 @@ class PipelineStage(PipelineStageBase):
178	return	190	return
179	for shared_param_info in self._shared_parameters:	191	for shared_param_info in self._shared_parameters:
180	param = shared_param_info.parameter	192	param = shared_param_info.parameter
		193	+ # On the meta-init -> FSDP-wrap -> load path the storage is still on
		194	+ # meta here; the trainer re-runs this after materialization, so skip
		195	+ # until it is real (broadcasting a meta tensor is invalid).
		196	+ if getattr(param, "is_meta", False):
		197	+ continue
181	shared_stage = shared_param_info.shared_stage	198	shared_stage = shared_param_info.shared_stage
182	group, group_ranks = self._init_shared_parameter_group(shared_stage)	199	group, group_ranks = self._init_shared_parameter_group(shared_stage)
183	shared_param_info.group = group	200	shared_param_info.group = group
		201	+ # DTensor dispatch whitelists DistCommBroadcast and unwraps DTensor
		202	+ # args to their local shards before invoking the backend collective.
184	platform.broadcast(param, group_ranks[0], group)	203	platform.broadcast(param, group_ranks[0], group)
185		204
186	def _global_rank(self, stage_index):	205	def _global_rank(self, stage_index):
	@@ -193,26 +212,65 @@ class PipelineStage(PipelineStageBase):
193		212
194	def _init_shared_parameter_group(self, shared_stage):	213	def _init_shared_parameter_group(self, shared_stage):
195	"""init group of shared parameter."""	214	"""init group of shared parameter."""
196	- group_ranks = []	215	+ group_ranks = [self._global_rank(stage) for stage in shared_stage]
197	- for stage in shared_stage:	216	+ # When the shared stages span the whole PP sub-mesh (the tied-embedding
198	- global_rank = self._global_rank(stage)	217	+ # pp=2 case), reuse the PP group: it is already created per dp line, so
199	- group_ranks.append(global_rank)	218	+ # this avoids a world-collective ``new_group`` that would deadlock under
		219	+ # FSDP where each dp line needs its own {stage0, stage1} group.
		220	+ if self.mesh is not None and len(group_ranks) == self.mesh.size():
		221	+ return self.mesh.get_group(), group_ranks
200	group = platform.create_group(group_ranks)	222	group = platform.create_group(group_ranks)
201	return group, group_ranks	223	return group, group_ranks
202		224
203	- def sync_shared_parameters_grad(self):	225	+ def sync_shared_parameters_grad(self) -> None:
204	"""sync shared parameters' grad with AllReduce."""	226	"""sync shared parameters' grad with AllReduce."""
205	if self._shared_parameters is None or not self._has_backward:	227	if self._shared_parameters is None or not self._has_backward:
206	return	228	return
207	for shared_param_info in self._shared_parameters:	229	for shared_param_info in self._shared_parameters:
208	param = shared_param_info.parameter	230	param = shared_param_info.parameter
		231	+ # Shared endpoints inherit the same tied/frozen configuration, so
		232	+ # ``requires_grad`` is group-wide. Frozen peers can all skip.
209	if not param.requires_grad:	233	if not param.requires_grad:
210	continue	234	continue
211	- grad = param.grad
212	group = shared_param_info.group	235	group = shared_param_info.group
		236	+ # ``group`` is assigned by ``_sync_shared_parameters`` once the
		237	+ # parameter is materialized off meta; if that handshake has not run
		238	+ # there is no peer to reduce with yet, so skip.
		239	+ if group is None:
		240	+ continue
		241	+
		242	+ grad = param.grad
		243	+ has_local_grad = grad is not None
		244	+ if has_local_grad:
		245	+ # Per-stage FSDP stores the reduced boundary-parameter grad as
		246	+ # a sharded DTensor. Reduce its local shard so matching DP
		247	+ # shards at the tied PP endpoints sum to one gradient.
		248	+ local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
		249	+ else:
		250	+ # Every member of a shared-parameter group must enter the same
		251	+ # collective even when this stage did not use the parameter in
		252	+ # backward. Derive the zero contribution from the live local
		253	+ # parameter shard so shape, dtype, and device match TP/FSDP.
		254	+ local_param = param.to_local() if isinstance(param, DTensor) else param
		255	+ local_grad = platform.full_like(local_param, 0)
		256	+
213	# platform.all_reduce expects group_info (with .group for Torch, or str for MindSpore)	257	# platform.all_reduce expects group_info (with .group for Torch, or str for MindSpore)
214	group_info = group if isinstance(group, str) else SimpleNamespace(group=group)	258	group_info = group if isinstance(group, str) else SimpleNamespace(group=group)
215	- platform.all_reduce(grad, group_info)	259	+ reduced_grad, handle = platform.all_reduce(local_grad, group_info, async_op=True)
		260	+ if handle is not None:
		261	+ handle.wait()
		262	+
		263	+ if has_local_grad:
		264	+ # Torch may return a contiguous communication tensor for a
		265	+ # non-contiguous input. Preserve the original grad object (and
		266	+ # DTensor layout) while copying the completed reduction back.
		267	+ if reduced_grad is not local_grad:
		268	+ platform.load_into_param(local_grad, reduced_grad)
		269	+ elif param.requires_grad:
		270	+ if isinstance(param, DTensor):
		271	+ param.grad = DTensor.from_local(reduced_grad, param.device_mesh, param.placements)
		272	+ else:
		273	+ param.grad = reduced_grad
216		274
217	def _check_src_stage(self, src_stage):	275	def _check_src_stage(self, src_stage):
218	"""check type for src_stage."""	276	"""check type for src_stage."""
	@@ -504,27 +562,47 @@ class PipelineStage(PipelineStageBase):
504	rg_infos = [info for info in self.args_recv_info[micro_index] if info.requires_grad]	562	rg_infos = [info for info in self.args_recv_info[micro_index] if info.requires_grad]
505	return [("isend", cur_out, info.global_rank) for cur_out, info in zip(out, rg_infos)]	563	return [("isend", cur_out, info.global_rank) for cur_out, info in zip(out, rg_infos)]
506		564
507	- def execute_reduce_grad(self):	565	+ def execute_reduce_grad(self) -> None:
508	- """Trigger FSDP post-backward gradient reduction and root hook for the stage submodule."""	566	+ """Reduce nested FSDP gradients and finalize every HSDP root in the stage."""
509	- if not isinstance(self.submodule, HSDPModule):	567	+ fsdp_schedulers = []
510	- return	568	+ seen_scheduler_ids = set()
511	- fsdp_module = self.submodule	569	+ for _, submod in platform.get_cells_and_names(self.submodule):
512	- fsdp_module.set_is_last_backward(True)
513	- fsdp_module.set_reshard_after_backward(True)
514	- fsdp_module.set_requires_gradient_sync(True)
515	-
516	- for _, submod in platform.get_cells_and_names(fsdp_module):
517	if not isinstance(submod, HSDPModule):	570	if not isinstance(submod, HSDPModule):
518	continue	571	continue
519	- sub_mod_state = submod.hsdp_scheduler.hsdp_state	572	+ scheduler = submod.hsdp_scheduler
520	- sub_mod_state.post_backward()	573	+ scheduler_id = id(scheduler)
521	- sub_mod_state.reduce_params()	574	+ if scheduler_id in seen_scheduler_ids:
		575	+ continue
		576	+ seen_scheduler_ids.add(scheduler_id)
		577	+ fsdp_schedulers.append(scheduler)
522		578
523	- # No public API exposes the root backward finalization; call the platform hook directly.	579	+ if not fsdp_schedulers:
524	- # force_reduce=True: the recv buffer's PostBackwardFunction has put the root into	580	+ return
525	- # scheduler_state==BACKWARD, so the natural gate would skip the final drain and the	581	+
526	- # last module's reduce-scatter would lag one optimizer step.	582	+ for scheduler in fsdp_schedulers:
527	- fsdp_module.hsdp_scheduler._root_backward_hook(force_reduce=True) # pylint: disable=protected-access	583	+ scheduler.scheduler_ctx.is_last_backward = True
		584	+ scheduler.set_reshard_after_backward(True)
		585	+ scheduler.set_requires_grad_sync(True)
		586	+
		587	+ for scheduler in fsdp_schedulers:
		588	+ scheduler.hsdp_state.post_backward()
		589	+ scheduler.hsdp_state.reduce_params()
		590	+
		591	+ # A plain pipeline-stage root may contain one or more independently
		592	+ # wrapped HSDP child roots. Finalize every runtime root so each context
		593	+ # clears its backward state; the platform hook also drains the global
		594	+ # fused reduction queues. A pre-forward fallback retains the historical
		595	+ # single-root behavior for callers that invoke this method directly.
		596	+ root_schedulers = [
		597	+ scheduler for scheduler in fsdp_schedulers
		598	+ if scheduler._is_root # pylint: disable=protected-access
		599	+ ]
		600	+ if not root_schedulers:
		601	+ root_schedulers = fsdp_schedulers[:1]
		602	+ for scheduler in root_schedulers:
		603	+ # No public API exposes root backward finalization. force_reduce=True
		604	+ # ensures a differentiable PP input cannot defer the final drain.
		605	+ scheduler._root_backward_hook(force_reduce=True) # pylint: disable=protected-access
528		606
529	def _build_padded_sens(self, micro_index):	607	def _build_padded_sens(self, micro_index):
530	"""Build an N-length sens list aligned with the forward output structure.	608	"""Build an N-length sens list aligned with the forward output structure.

Mhyper_parallel/core/shard/_op_dispatch.py+41-1

	@@ -566,6 +566,37 @@ class OpDispatcher:
566	skip_dispatch = get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()	566	skip_dispatch = get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()
567	return op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS or skip_dispatch	567	return op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS or skip_dispatch
568		568
		569	+ @staticmethod
		570	+ def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None:
		571	+ """Reject local in-place add/sub when Partial contributions need gating."""
		572	+ if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args:
		573	+ return
		574	+ first = args[0]
		575	+ if len(args) >= 2:
		576	+ second = args[1]
		577	+ elif "other" in kwargs:
		578	+ second = kwargs["other"]
		579	+ else:
		580	+ return
		581	+ if not isinstance(first, DTensor):
		582	+ return
		583	+ mesh_ndim = len(first.layout.partial)
		584	+ first_partial = tuple(first.layout.partial)
		585	+ if isinstance(second, DTensor):
		586	+ second_partial = tuple(second.layout.partial)
		587	+ if len(second_partial) != mesh_ndim:
		588	+ raise ValueError(
		589	+ f"For {op_name}, in-place input mesh dimensions must match, "
		590	+ f"but got {mesh_ndim} and {len(second_partial)}."
		591	+ )
		592	+ else:
		593	+ second_partial = (None,) * mesh_ndim
		594	+ if first_partial != second_partial:
		595	+ raise ValueError(
		596	+ f"For {op_name}, input Partial placements must be identical for "
		597	+ f"local in-place execution, but got {first_partial} and {second_partial}."
		598	+ )
		599	+
569	def _should_dispatch_loss_parallel(self, op_name: str) -> bool:	600	def _should_dispatch_loss_parallel(self, op_name: str) -> bool:
570	"""Check if should dispatch through loss_parallel path.	601	"""Check if should dispatch through loss_parallel path.
571		602
	@@ -791,7 +822,15 @@ class OpDispatcher:
791	f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. "	822	f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. "
792	f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)."	823	f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)."
793	)	824	)
794	- return self._dispatch_new(op_call, distribute_op, packed_call, result)	825	+ output = self._dispatch_new(op_call, distribute_op, packed_call, result)
		826	+ return self._restore_inplace_dtensor_result(op_name, args, output)
		827	+
		828	+ @staticmethod
		829	+ def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any:
		830	+ """Return the original DTensor wrapper after a local in-place operation."""
		831	+ if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor):
		832	+ return args[0]
		833	+ return output
795		834
796	@staticmethod	835	@staticmethod
797	def _dispatch_new(func, distribute_op, packed_call, result) -> Tensor:	836	def _dispatch_new(func, distribute_op, packed_call, result) -> Tensor:
	@@ -852,6 +891,7 @@ class OpDispatcher:
852	op_name = platform.get_op_name(op_call)	891	op_name = platform.get_op_name(op_call)
853		892
854	if self._should_bypass_dispatch(op_name):	893	if self._should_bypass_dispatch(op_name):
		894	+ self._validate_inplace_partial_inputs(op_name, args, kwargs)
855	result = op_call(self._unwrap_args(args), *self._unwrap_kwargs(kwargs))	895	result = op_call(self._unwrap_args(args), *self._unwrap_kwargs(kwargs))
856	if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):	896	if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):
857	return args[0]	897	return args[0]

Mhyper_parallel/core/shard/ops/parallel_elementwise.py+143-17

	@@ -22,6 +22,59 @@ from typing import Callable, Optional, Tuple
22	from .parallel_ops import DistributedOp	22	from .parallel_ops import DistributedOp
23		23
24		24
		25	+_INPLACE_ELEMENTWISE_OPS = frozenset({
		26	+ "add_", "sub_", "InplaceAddExt", "InplaceSubExt",
		27	+})
		28	+
		29	+
		30	+def _partial_signature(layout, mesh_ndim: int) -> tuple:
		31	+ """Return one Partial entry per output mesh dimension."""
		32	+ if layout is None:
		33	+ return (None,) * mesh_ndim
		34	+ partial = tuple(layout.partial)
		35	+ if len(partial) != mesh_ndim:
		36	+ raise ValueError(
		37	+ f"Input and output mesh dimensions must match, but got "
		38	+ f"input={len(partial)} and output={mesh_ndim}."
		39	+ )
		40	+ return partial
		41	+
		42	+
		43	+def _contributes_to_partial_output(input_layout, output_layout) -> bool:
		44	+ """Return whether this rank contributes an input to a Partial output.
		45	+
		46	+ An input that is not Partial on one of the output's Partial axes is
		47	+ replicated along that axis. Only coordinate zero may contribute that
		48	+ replicated value, otherwise the eventual reduction would count it once
		49	+ per rank on the added axis.
		50	+ """
		51	+ output_partial = tuple(output_layout.partial)
		52	+ input_partial = _partial_signature(input_layout, len(output_partial))
		53	+ for mesh_dim, output_partial_type in enumerate(output_partial):
		54	+ if (
		55	+ output_partial_type is not None
		56	+ and input_partial[mesh_dim] is None
		57	+ and output_layout.mesh.get_local_rank(mesh_dim) != 0
		58	+ ):
		59	+ return False
		60	+ return True
		61	+
		62	+
		63	+def _zero_contribution(value):
		64	+ """Create a strict zero while retaining a floating tensor's grad edge."""
		65	+ is_complex = getattr(value, "is_complex", None)
		66	+ if callable(is_complex) and is_complex():
		67	+ return value * 0
		68	+ if hasattr(value, "clamp"):
		69	+ # ``value * 0`` turns +/-Inf into NaN. Clamp first so the forward value
		70	+ # is finite, then multiply by zero to make the derivative zero even at
		71	+ # the clamp boundary.
		72	+ return value.clamp(0, 0) * 0
		73	+ if isinstance(value, (bool, int, float, complex)):
		74	+ return type(value)(0)
		75	+ return value * 0
		76	+
		77	+
25	def _unwrap_local_value(value):	78	def _unwrap_local_value(value):
26	"""Convert DTensor-like values to local tensors while preserving containers."""	79	"""Convert DTensor-like values to local tensors while preserving containers."""
27	if hasattr(value, "_layout"):	80	if hasattr(value, "_layout"):
	@@ -604,6 +657,48 @@ class AddDistributedOp(ElementWiseWithPartialDistributedOp):
604	which is useful for operations like gradient accumulation where partial	657	which is useful for operations like gradient accumulation where partial
605	results need to be preserved through the computation graph.	658	results need to be preserved through the computation graph.
606	"""	659	"""
		660	+
		661	+ @staticmethod
		662	+ def _canonicalize_binary_operands(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
		663	+ """Move semantic ``input`` and ``other`` operands before option values."""
		664	+ local_kwargs = dict(kwargs)
		665	+ if args:
		666	+ if "input" in local_kwargs:
		667	+ raise ValueError("Binary elementwise input was provided both positionally and by keyword.")
		668	+ lhs = args[0]
		669	+ trailing_args = args[1:]
		670	+ elif "input" in local_kwargs:
		671	+ lhs = local_kwargs.pop("input")
		672	+ trailing_args = ()
		673	+ else:
		674	+ return args, local_kwargs
		675	+
		676	+ if trailing_args:
		677	+ if "other" in local_kwargs:
		678	+ raise ValueError("Binary elementwise other was provided both positionally and by keyword.")
		679	+ rhs = trailing_args[0]
		680	+ trailing_args = trailing_args[1:]
		681	+ elif "other" in local_kwargs:
		682	+ rhs = local_kwargs.pop("other")
		683	+ else:
		684	+ return args, local_kwargs
		685	+ return (lhs, rhs, *trailing_args), local_kwargs
		686	+
		687	+ def preprocess(self, args: tuple, kwargs: dict) -> tuple:
		688	+ """Canonicalize binary operands before local unwrapping and layout caching.
		689	+
		690	+ Args:
		691	+ args: Positional backend arguments.
		692	+ kwargs: Keyword backend arguments, potentially including
		693	+ ``input``, ``other``, and options such as ``alpha``.
		694	+
		695	+ Returns:
		696	+ Local arguments, local keyword arguments, and layout cache values
		697	+ with the semantic left and right operands in the first two slots.
		698	+ """
		699	+ canonical_args, canonical_kwargs = self._canonicalize_binary_operands(args, kwargs)
		700	+ return super().preprocess(canonical_args, canonical_kwargs)
		701	+
607	def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221	702	def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
608	"""	703	"""
609	Infer output layout for Add operator.	704	Infer output layout for Add operator.
	@@ -622,6 +717,22 @@ class AddDistributedOp(ElementWiseWithPartialDistributedOp):
622	Raises:	717	Raises:
623	ValueError: If propagated Partial status is not "sum" or None.	718	ValueError: If propagated Partial status is not "sum" or None.
624	"""	719	"""
		720	+ layouts = tuple(cache_values[:-1])
		721	+ if self.op_name in _INPLACE_ELEMENTWISE_OPS and len(layouts) >= 2:
		722	+ mesh_ndim = max(
		723	+ (len(layout.partial) for layout in layouts[:2] if layout is not None),
		724	+ default=0,
		725	+ )
		726	+ partial_signatures = tuple(
		727	+ _partial_signature(layout, mesh_ndim)
		728	+ for layout in layouts[:2]
		729	+ )
		730	+ if partial_signatures[0] != partial_signatures[1]:
		731	+ raise ValueError(
		732	+ f"For {self.op_name}, input Partial placements should be identical "
		733	+ f"for in-place execution, but got {partial_signatures}."
		734	+ )
		735	+
625	infer_result = super().infer_layout(cache_values)	736	infer_result = super().infer_layout(cache_values)
626	if infer_result is None:	737	if infer_result is None:
627	return infer_result	738	return infer_result
	@@ -638,31 +749,46 @@ class AddDistributedOp(ElementWiseWithPartialDistributedOp):
638		749
639	def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221	750	def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
640	cache_values: list) -> Optional[Callable]:	751	cache_values: list) -> Optional[Callable]:
641	- """	752	+ """Build a rank-local implementation that avoids repeated values.
642	- Get expand implementation for the operator	753	+
		754	+ Args:
		755	+ func: Local backend add or subtract implementation.
		756	+ infer_result: Inferred output layout and auxiliary result.
		757	+ cache_values: Input layouts followed by input shapes.
		758	+
		759	+ Returns:
		760	+ A contribution-aware local implementation, or ``None`` when both
		761	+ inputs already have identical Partial placements.
643	"""	762	"""
644	layouts = tuple(cache_values[:-1])	763	layouts = tuple(cache_values[:-1])
645	x1_layout = layouts[0]	764	x1_layout = layouts[0]
646	x2_layout = layouts[1]	765	x2_layout = layouts[1]
647	- x1_partial = x1_layout.is_partial() if x1_layout is not None else False	766	+ output_layout = infer_result[0][0]
648	- x2_partial = x2_layout.is_partial() if x2_layout is not None else False	767	+ output_mesh_ndim = len(output_layout.partial)
		768	+ x1_partial = _partial_signature(x1_layout, output_mesh_ndim)
		769	+ x2_partial = _partial_signature(x2_layout, output_mesh_ndim)
649		770
650	if x1_partial == x2_partial:	771	if x1_partial == x2_partial:
651	return None	772	return None
652		773
653	- output_layout = infer_result[0][0]	774	+ x1_contributes = _contributes_to_partial_output(x1_layout, output_layout)
654	- scaling_factor = 1	775	+ x2_contributes = _contributes_to_partial_output(x2_layout, output_layout)
655	- for i, partial_type in enumerate(output_layout.partial):
656	- if partial_type == "sum":
657	- scaling_factor *= output_layout.mesh_shape[i]
658		776
659	- # use expand_impl only when one of x1 and x2 is with partial placement.	777	+ # ``args``/``*kwargs`` forward add-specific extras such as ``alpha``.
660	- def _expand_impl1(x1, x2, args, *kwargs):	778	+ # A strict zero retains a valid zero-gradient autograd edge.
661	- add_out = func(x1 / scaling_factor, x2, args, *kwargs)	779	+ def _expand_impl(x1, x2, args, *kwargs):
662	- return add_out	780	+ local_x1 = x1 if x1_contributes else _zero_contribution(x1)
		781	+ local_x2 = x2 if x2_contributes else _zero_contribution(x2)
		782	+ return func(local_x1, local_x2, args, *kwargs)
663		783
664	- def _expand_impl2(x1, x2, args, *kwargs):	784	+ return _expand_impl
665	- add_out = func(x1, x2 / scaling_factor, args, *kwargs)
666	- return add_out
667		785
668	- return _expand_impl1 if not x1_partial else _expand_impl2	786	+
		787	+class SubDistributedOp(AddDistributedOp):
		788	+ """Distributed subtraction with correct replicated-minus-partial signs.
		789	+
		790	+ When the first operand is replicated and the second is Partial(sum), only
		791	+ the first Partial rank may contribute the replicated value. Every other
		792	+ rank must contribute ``0 - x2_local`` so reducing the output reconstructs
		793	+ ``x1 - sum(x2_local)``.
		794	+ """

Mhyper_parallel/core/shard/ops/yaml/element_wise_ops_with_shape.yaml+2-2

	@@ -39,7 +39,7 @@ Add:
39		39
40	Sub:	40	Sub:
41	dist_op_name: _sub_dist_op	41	dist_op_name: _sub_dist_op
42	- distributed_op_class: AddDistributedOp	42	+ distributed_op_class: SubDistributedOp
43		43
44	Mul:	44	Mul:
45	dist_op_name: _mul_dist_op	45	dist_op_name: _mul_dist_op
	@@ -57,7 +57,7 @@ div_:
57		57
58	SubExt:	58	SubExt:
59	dist_op_name: _sub_ext_dist_op	59	dist_op_name: _sub_ext_dist_op
60	- distributed_op_class: AddDistributedOp	60	+ distributed_op_class: SubDistributedOp
61		61
62	isnan:	62	isnan:
63	dist_op_name: _torch_isnan_dist_op	63	dist_op_name: _torch_isnan_dist_op

Mhyper_parallel/core/shard/ops/yaml/torch_element_wise.yaml+2-1

	@@ -7,6 +7,7 @@ relu:
7		7
8	__rsub__:	8	__rsub__:
9	dist_op_name: _torch___sub___dist_op	9	dist_op_name: _torch___sub___dist_op
		10	+ distributed_op_class: SubDistributedOp
10		11
11	clone:	12	clone:
12	dist_op_name: _torch_clone_dist_op	13	dist_op_name: _torch_clone_dist_op
	@@ -67,4 +68,4 @@ zeros_like:
67	dist_op_name: _torch_zeros_like_dist_op	68	dist_op_name: _torch_zeros_like_dist_op
68		69
69	empty_like:	70	empty_like:
70	- dist_op_name: _torch_empty_like_dist_op	71	+ dist_op_name: _torch_empty_like_dist_op

Mhyper_parallel/core/shard/ops/yaml/torch_element_wise_with_shape.yaml+1-0

	@@ -12,6 +12,7 @@ add_:
12		12
13	sub:	13	sub:
14	dist_op_name: _torch_sub_dist_op	14	dist_op_name: _torch_sub_dist_op
		15	+ distributed_op_class: SubDistributedOp
15		16
16	sub_:	17	sub_:
17	dist_op_name: _torch_sub__dist_op	18	dist_op_name: _torch_sub__dist_op

Mhyper_parallel/core/tensor_parallel/style.py+62-4

	@@ -38,6 +38,14 @@ from hyper_parallel.platform import get_platform
38	platform = get_platform()	38	platform = get_platform()
39	Module = platform.Module	39	Module = platform.Module
40		40
		41	+
		42	+def _src_data_rank_for_tensor(tensor: Any, src_data_rank: Optional[int]) -> Optional[int]:
		43	+ """Disable source-rank communication while sharding meta tensors."""
		44	+ if getattr(tensor, "is_meta", False):
		45	+ return None
		46	+ return src_data_rank
		47	+
		48	+
41	__all__ = [	49	__all__ = [
42	"ParallelStyle",	50	"ParallelStyle",
43	"ColwiseParallel",	51	"ColwiseParallel",
	@@ -158,7 +166,12 @@ class ColwiseParallel(ParallelStyle):
158	continue	166	continue
159	src = _distribute_module_param_source(param)	167	src = _distribute_module_param_source(param)
160	requires_grad = bool(getattr(param, "requires_grad", True))	168	requires_grad = bool(getattr(param, "requires_grad", True))
161	- dt = distribute_tensor(src, device_mesh, [Shard(0)], src_data_rank=self.src_data_rank)	169	+ dt = distribute_tensor(
		170	+ src,
		171	+ device_mesh,
		172	+ [Shard(0)],
		173	+ src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank),
		174	+ )
162	new_param = _distribute_module_new_parameter(key, dt, requires_grad)	175	new_param = _distribute_module_new_parameter(key, dt, requires_grad)
163	_distribute_module_set_param(module, key, new_param)	176	_distribute_module_set_param(module, key, new_param)
164		177
	@@ -169,7 +182,12 @@ class ColwiseParallel(ParallelStyle):
169	continue	182	continue
170	src = _distribute_module_param_source(param)	183	src = _distribute_module_param_source(param)
171	requires_grad = bool(getattr(param, "requires_grad", True))	184	requires_grad = bool(getattr(param, "requires_grad", True))
172	- dt = distribute_tensor(src, device_mesh, [Shard(1)], src_data_rank=self.src_data_rank)	185	+ dt = distribute_tensor(
		186	+ src,
		187	+ device_mesh,
		188	+ [Shard(1)],
		189	+ src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank),
		190	+ )
173	new_param = _distribute_module_new_parameter(key, dt, requires_grad)	191	new_param = _distribute_module_new_parameter(key, dt, requires_grad)
174	_distribute_module_set_param(module, key, new_param)	192	_distribute_module_set_param(module, key, new_param)
175		193
	@@ -254,6 +272,10 @@ class RowwiseParallel(ParallelStyle):
254	output_layouts (Placement, optional):	272	output_layouts (Placement, optional):
255	Desired DTensor layout of the module output. Defaults to	273	Desired DTensor layout of the module output. Defaults to
256	``Replicate()`` (all-reduce / reduce-scatter from partial).	274	``Replicate()`` (all-reduce / reduce-scatter from partial).
		275	+ reduce_dtype (dtype, optional):
		276	+ Floating-point dtype used for reducing Partial outputs; the reduced
		277	+ output remains in this dtype. Defaults to ``None`` (reduce in the
		278	+ output dtype).
257	use_local_output (bool, optional):	279	use_local_output (bool, optional):
258	If ``True`` (default), convert the output DTensor back to a local	280	If ``True`` (default), convert the output DTensor back to a local
259	tensor via ``to_local()``.	281	tensor via ``to_local()``.
	@@ -273,12 +295,14 @@ class RowwiseParallel(ParallelStyle):
273	*,	295	*,
274	input_layouts: Optional[Placement] = None,	296	input_layouts: Optional[Placement] = None,
275	output_layouts: Optional[Placement] = None,	297	output_layouts: Optional[Placement] = None,
		298	+ reduce_dtype: Optional[Any] = None,
276	use_local_output: bool = True,	299	use_local_output: bool = True,
277	) -> None:	300	) -> None:
278	super().__init__()	301	super().__init__()
279	self.input_layouts: Tuple[Placement, ...] = (input_layouts or Shard(-1),)	302	self.input_layouts: Tuple[Placement, ...] = (input_layouts or Shard(-1),)
280	self.output_layouts: Tuple[Placement, ...] = (output_layouts or Replicate(),)	303	self.output_layouts: Tuple[Placement, ...] = (output_layouts or Replicate(),)
281	self.desired_input_layouts: Tuple[Placement, ...] = (Shard(-1),)	304	self.desired_input_layouts: Tuple[Placement, ...] = (Shard(-1),)
		305	+ self.reduce_dtype = reduce_dtype
282	self.use_local_output = use_local_output	306	self.use_local_output = use_local_output
283		307
284	def __repr__(self) -> str:	308	def __repr__(self) -> str:
	@@ -286,6 +310,7 @@ class RowwiseParallel(ParallelStyle):
286	f"{self.__class__.__name__}("	310	f"{self.__class__.__name__}("
287	f"input_layouts={self.input_layouts}, "	311	f"input_layouts={self.input_layouts}, "
288	f"output_layouts={self.output_layouts}, "	312	f"output_layouts={self.output_layouts}, "
		313	+ f"reduce_dtype={self.reduce_dtype}, "
289	f"use_local_output={self.use_local_output})"	314	f"use_local_output={self.use_local_output})"
290	)	315	)
291		316
	@@ -318,7 +343,12 @@ class RowwiseParallel(ParallelStyle):
318	src = _distribute_module_param_source(param)	343	src = _distribute_module_param_source(param)
319	requires_grad = bool(getattr(param, "requires_grad", True))	344	requires_grad = bool(getattr(param, "requires_grad", True))
320	placement = [Shard(1)] if key == "weight" else [Replicate()]	345	placement = [Shard(1)] if key == "weight" else [Replicate()]
321	- dt = distribute_tensor(src, device_mesh, placement, src_data_rank=self.src_data_rank)	346	+ dt = distribute_tensor(
		347	+ src,
		348	+ device_mesh,
		349	+ placement,
		350	+ src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank),
		351	+ )
322	new_param = _distribute_module_new_parameter(key, dt, requires_grad)	352	new_param = _distribute_module_new_parameter(key, dt, requires_grad)
323	_distribute_module_set_param(module, key, new_param)	353	_distribute_module_set_param(module, key, new_param)
324		354
	@@ -329,7 +359,12 @@ class RowwiseParallel(ParallelStyle):
329	continue	359	continue
330	src = _distribute_module_param_source(param)	360	src = _distribute_module_param_source(param)
331	requires_grad = bool(getattr(param, "requires_grad", True))	361	requires_grad = bool(getattr(param, "requires_grad", True))
332	- dt = distribute_tensor(src, device_mesh, [Shard(0)], src_data_rank=self.src_data_rank)	362	+ dt = distribute_tensor(
		363	+ src,
		364	+ device_mesh,
		365	+ [Shard(0)],
		366	+ src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank),
		367	+ )
333	new_param = _distribute_module_new_parameter(key, dt, requires_grad)	368	new_param = _distribute_module_new_parameter(key, dt, requires_grad)
334	_distribute_module_set_param(module, key, new_param)	369	_distribute_module_set_param(module, key, new_param)
335		370
	@@ -340,6 +375,7 @@ class RowwiseParallel(ParallelStyle):
340	outputs: Any,	375	outputs: Any,
341	device_mesh: DeviceMesh,	376	device_mesh: DeviceMesh,
342	module: Optional[Module] = None,	377	module: Optional[Module] = None,
		378	+ reduce_dtype: Optional[Any] = None,
343	) -> Any:	379	) -> Any:
344	"""Redistribute partial output and optionally convert to local."""	380	"""Redistribute partial output and optionally convert to local."""
345	if not isinstance(outputs, DTensor):	381	if not isinstance(outputs, DTensor):
	@@ -352,6 +388,9 @@ class RowwiseParallel(ParallelStyle):
352	"RowwiseParallel expects a DTensor from Linear outputs; "	388	"RowwiseParallel expects a DTensor from Linear outputs; "
353	f"got {type(outputs)}. If this is an unsupported module, extend I/O hooks."	389	f"got {type(outputs)}. If this is an unsupported module, extend I/O hooks."
354	)	390	)
		391	+ if reduce_dtype is not None and tuple(outputs.placements) != tuple(output_layouts):
		392	+ local_output = platform.cast_fp_tensor(reduce_dtype, outputs.to_local())
		393	+ outputs = DTensor.from_local(local_output, outputs.device_mesh, outputs.placements)
355	if tuple(outputs.placements) != tuple(output_layouts):	394	if tuple(outputs.placements) != tuple(output_layouts):
356	outputs = outputs.redistribute(device_mesh, output_layouts)	395	outputs = outputs.redistribute(device_mesh, output_layouts)
357	if use_local_output:	396	if use_local_output:
	@@ -403,6 +442,7 @@ class RowwiseParallel(ParallelStyle):
403	forward_outputs,	442	forward_outputs,
404	device_mesh,	443	device_mesh,
405	forward_module,	444	forward_module,
		445	+ self.reduce_dtype,
406	)	446	)
407		447
408	return distribute_module(	448	return distribute_module(
	@@ -675,6 +715,8 @@ class PrepareModuleOutput(ParallelStyle):
675	Keyword Args:	715	Keyword Args:
676	output_layouts: Current or assumed placement per output tensor.	716	output_layouts: Current or assumed placement per output tensor.
677	desired_output_layouts: Target placements; length must match ``output_layouts``.	717	desired_output_layouts: Target placements; length must match ``output_layouts``.
		718	+ reduce_dtype: Floating-point dtype used before reducing Partial outputs;
		719	+ the reduced output remains in this dtype.
678	use_local_output: If ``True`` (default), return local shards after redistribution.	720	use_local_output: If ``True`` (default), return local shards after redistribution.
679	"""	721	"""
680		722
	@@ -683,6 +725,7 @@ class PrepareModuleOutput(ParallelStyle):
683	*,	725	*,
684	output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],	726	output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],
685	desired_output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],	727	desired_output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],
		728	+ reduce_dtype: Optional[Any] = None,
686	use_local_output: bool = True,	729	use_local_output: bool = True,
687	) -> None:	730	) -> None:
688	super().__init__()	731	super().__init__()
	@@ -694,6 +737,7 @@ class PrepareModuleOutput(ParallelStyle):
694	if isinstance(desired_output_layouts, Placement)	737	if isinstance(desired_output_layouts, Placement)
695	else desired_output_layouts	738	else desired_output_layouts
696	)	739	)
		740	+ self.reduce_dtype = reduce_dtype
697	self.use_local_output = use_local_output	741	self.use_local_output = use_local_output
698	if len(self.output_layouts) != len(self.desired_output_layouts):	742	if len(self.output_layouts) != len(self.desired_output_layouts):
699	raise AssertionError(	743	raise AssertionError(
	@@ -715,6 +759,16 @@ class PrepareModuleOutput(ParallelStyle):
715	dt_out = out	759	dt_out = out
716	else:	760	else:
717	dt_out = DTensor.from_local(out, device_mesh, (out_layout,))	761	dt_out = DTensor.from_local(out, device_mesh, (out_layout,))
		762	+ has_partial_output = any(
		763	+ placement.is_partial() for placement in dt_out.placements
		764	+ )
		765	+ if (
		766	+ self.reduce_dtype is not None
		767	+ and has_partial_output
		768	+ and out_layout != desired_out_layout
		769	+ ):
		770	+ local_out = platform.cast_fp_tensor(self.reduce_dtype, dt_out.to_local())
		771	+ dt_out = DTensor.from_local(local_out, dt_out.device_mesh, dt_out.placements)
718	if out_layout != desired_out_layout:	772	if out_layout != desired_out_layout:
719	dt_out = dt_out.redistribute(device_mesh, (desired_out_layout,))	773	dt_out = dt_out.redistribute(device_mesh, (desired_out_layout,))
720	prepared_outputs.append(	774	prepared_outputs.append(
	@@ -739,6 +793,7 @@ class PrepareModuleOutput(ParallelStyle):
739	f"{self.__class__.__name__}("	793	f"{self.__class__.__name__}("
740	f"output_layouts={self.output_layouts}, "	794	f"output_layouts={self.output_layouts}, "
741	f"desired_output_layouts={self.desired_output_layouts}, "	795	f"desired_output_layouts={self.desired_output_layouts}, "
		796	+ f"reduce_dtype={self.reduce_dtype}, "
742	f"use_local_output={self.use_local_output})"	797	f"use_local_output={self.use_local_output})"
743	)	798	)
744		799
	@@ -762,6 +817,7 @@ class PrepareModuleInputOutput(ParallelStyle):
762	use_local_input: bool = False,	817	use_local_input: bool = False,
763	output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],	818	output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],
764	desired_output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],	819	desired_output_layouts: Union[Placement, Tuple[Optional[Placement], ...]],
		820	+ reduce_dtype: Optional[Any] = None,
765	use_local_output: bool = True,	821	use_local_output: bool = True,
766	) -> None:	822	) -> None:
767	super().__init__()	823	super().__init__()
	@@ -775,6 +831,7 @@ class PrepareModuleInputOutput(ParallelStyle):
775	self.prepare_module_output = PrepareModuleOutput(	831	self.prepare_module_output = PrepareModuleOutput(
776	output_layouts=output_layouts,	832	output_layouts=output_layouts,
777	desired_output_layouts=desired_output_layouts,	833	desired_output_layouts=desired_output_layouts,
		834	+ reduce_dtype=reduce_dtype,
778	use_local_output=use_local_output,	835	use_local_output=use_local_output,
779	)	836	)
780		837
	@@ -795,6 +852,7 @@ class PrepareModuleInputOutput(ParallelStyle):
795	f"use_local_input={p_in.use_local_output}, "	852	f"use_local_input={p_in.use_local_output}, "
796	f"output_layouts={p_out.output_layouts}, "	853	f"output_layouts={p_out.output_layouts}, "
797	f"desired_output_layouts={p_out.desired_output_layouts}, "	854	f"desired_output_layouts={p_out.desired_output_layouts}, "
		855	+ f"reduce_dtype={p_out.reduce_dtype}, "
798	f"use_local_output={p_out.use_local_output})"	856	f"use_local_output={p_out.use_local_output})"
799	)	857	)
800		858

Ahyper_parallel/data/__init__.py+64-0

@@ -0,0 +1,64 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""HyperParallel dataset module for the unified ``data.type`` registry.
	16	+
	17	+Public API:
	18	+
	19	+- :func:`build_dataset`: registry-driven dispatch on ``data.type``.
	20	+- :data:`DATASET_REGISTRY`: registry instance; register custom formats
	21	+ with ``@DATASET_REGISTRY.register("<name>")``.
	22	+- :class:`DummyDataset`, :class:`DummyVLDataset`, :class:`TokenizedDataset`,
	23	+ :class:`PresetPtDataset`: first-party dataset classes.
	24	+- :class:`IndexedDataset`, :class:`IndexedDatasetBuilder`,
	25	+ :class:`GPTDataset`, :class:`BlendableDataset`: Megatron support.
	26	+
	27	+Importing this package registers every first-party builder
	28	+(dummy, vl_dummy, hf_datasets, json_file, preset_pt, megatron).
	29	+External plugins can register additional builders by importing their
	30	+modules before the trainer is constructed.
	31	+"""
	32	+from hyper_parallel.data.dummy import DummyDataset
	33	+from hyper_parallel.data.hf import TokenizedDataset
	34	+from hyper_parallel.data.megatron import (
	35	+ BlendableDataset,
	36	+ GPTDataset,
	37	+ IndexedDataset,
	38	+ IndexedDatasetBuilder,
	39	+)
	40	+from hyper_parallel.data.preset_pt import PresetPtDataset
	41	+from hyper_parallel.data.registry import DATASET_REGISTRY, build_dataset
	42	+from hyper_parallel.data.vl_dummy import DummyVLDataset
	43	+
	44	+# Each module registers a builder when imported.
	45	+from hyper_parallel.data import ( # noqa: F401 # registers builders
	46	+ dummy,
	47	+ hf,
	48	+ preset_pt,
	49	+ vl_dummy,
	50	+)
	51	+
	52	+
	53	+__all__ = [
	54	+ "BlendableDataset",
	55	+ "DATASET_REGISTRY",
	56	+ "DummyDataset",
	57	+ "DummyVLDataset",
	58	+ "GPTDataset",
	59	+ "IndexedDataset",
	60	+ "IndexedDatasetBuilder",
	61	+ "PresetPtDataset",
	62	+ "TokenizedDataset",
	63	+ "build_dataset",
	64	+]

Ahyper_parallel/data/dummy.py+71-0

@@ -0,0 +1,71 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Deterministic dummy LM dataset.
	16	+
	17	+Each sample's tokens are derived from ``base_seed + idx`` so the same global
	18	+sample index always yields the same content, independent of DP topology or
	19	+sampler shuffling.
	20	+"""
	21	+import logging
	22	+from typing import Any, Dict
	23	+
	24	+import torch
	25	+from torch.utils.data import Dataset
	26	+
	27	+from hyper_parallel.data.registry import DATASET_REGISTRY
	28	+
	29	+
	30	+logger = logging.getLogger(__name__)
	31	+
	32	+
	33	+class DummyDataset(Dataset):
	34	+ """Deterministic random-token dataset.
	35	+
	36	+ Args:
	37	+ num_samples: Total dataset length.
	38	+ seq_length: Tokens per sample.
	39	+ vocab: Token range upper bound (exclusive).
	40	+ base_seed: Seed offset; each sample uses ``base_seed + idx``.
	41	+ """
	42	+
	43	+ def __init__(self, num_samples: int, seq_length: int, vocab: int, base_seed: int = 42) -> None:
	44	+ self.num_samples = int(num_samples)
	45	+ self.seq_length = int(seq_length)
	46	+ self.vocab = int(vocab)
	47	+ self.base_seed = int(base_seed)
	48	+
	49	+ def __len__(self) -> int:
	50	+ return self.num_samples
	51	+
	52	+ def __getitem__(self, idx: int) -> Dict[str, Any]:
	53	+ g = torch.Generator().manual_seed(self.base_seed + idx)
	54	+ input_ids = torch.randint(0, self.vocab, (self.seq_length,), generator=g)
	55	+ return {"input_ids": input_ids, "labels": input_ids.clone()}
	56	+
	57	+
	58	+@DATASET_REGISTRY.register("dummy")
	59	+def build_dummy(, base: Any, args: Any, *_: Any) -> DummyDataset:
	60	+ """Build the dummy LM dataset from trainer config.
	61	+
	62	+ Picks ``vocab`` from ``model.config.vocab_size`` when available, else
	63	+ defaults to 32000 (matches the pre-refactor behaviour in BaseTrainer).
	64	+ """
	65	+ seq_len = args.data.max_seq_len
	66	+ vocab_size = base.model.config.vocab_size
	67	+ base_seed = args.train.seed
	68	+ total_samples = base.state.max_steps * args.train.global_batch_size
	69	+ ds = DummyDataset(total_samples, seq_len, vocab_size, base_seed=base_seed)
	70	+ logger.info("Dummy dataset created: %d samples, seq_len=%d", total_samples, seq_len)
	71	+ return ds

Ahyper_parallel/data/hf.py+141-0

@@ -0,0 +1,141 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""HuggingFace ``datasets`` builders (hub / local JSON).
	16	+
	17	+Two ``data.type`` aliases land here:
	18	+
	19	+- ``hf_datasets``: ``load_dataset(train_path, split="train")`` (hub or
	20	+ local arrow directory). Tokenises with the spec-provided transform.
	21	+- ``json_file``: ``load_dataset("json", data_files=train_path, split="train")``
	22	+ — covers Alpaca-style ``instruction / input / output`` files.
	23	+
	24	+The wrapped :class:`TokenizedDataset` returns plain tensor rows so the
	25	+trainer's standard padding collator can stack them. ``max_steps``
	26	+is auto-clipped to ``num_train_epochs * floor(len / global_batch_size)``
	27	+so the dataloader never runs dry mid-epoch.
	28	+"""
	29	+import logging
	30	+from typing import Any
	31	+
	32	+import torch
	33	+from torch.utils.data import Dataset
	34	+
	35	+from hyper_parallel.data.registry import DATASET_REGISTRY
	36	+
	37	+
	38	+logger = logging.getLogger(__name__)
	39	+
	40	+
	41	+class TokenizedDataset(Dataset):
	42	+ """Lightweight ``torch.utils.data.Dataset`` view over an HF dataset.
	43	+
	44	+ Each row is converted to ``torch.long`` tensors on demand — the HF
	45	+ table stays on disk / memory in arrow format until the dataloader
	46	+ actually pulls a sample.
	47	+ """
	48	+
	49	+ def __init__(self, hf_ds: Any) -> None:
	50	+ self.data = hf_ds
	51	+
	52	+ def __len__(self) -> int:
	53	+ return len(self.data)
	54	+
	55	+ def __getitem__(self, idx: int):
	56	+ item = self.data[idx]
	57	+ return {
	58	+ "input_ids": torch.tensor(item["input_ids"], dtype=torch.long),
	59	+ "labels": torch.tensor(item["labels"], dtype=torch.long),
	60	+ }
	61	+
	62	+
	63	+def _load_raw(args: Any, data_type: str) -> Any:
	64	+ """Run the appropriate ``load_dataset`` call for ``data_type``."""
	65	+ from datasets import load_dataset # pylint: disable=C0415 # optional dep
	66	+
	67	+ train_path = args.data.train_path
	68	+ if not train_path:
	69	+ raise ValueError(f"data.train_path is required when data.type='{data_type}'")
	70	+
	71	+ if data_type == "json_file":
	72	+ return load_dataset("json", data_files=train_path, split="train")
	73	+ subset = args.data.subset
	74	+ if subset:
	75	+ return load_dataset(train_path, subset, split="train")
	76	+ return load_dataset(train_path, split="train")
	77	+
	78	+
	79	+def _maybe_truncate(ds: Any, args: Any) -> Any:
	80	+ """Apply ``data.train_size`` if it shrinks the dataset."""
	81	+ train_size = args.data.train_size
	82	+ if train_size and train_size < len(ds):
	83	+ ds = ds.select(range(train_size))
	84	+ logger.info("Dataset truncated to %d samples", train_size)
	85	+ return ds
	86	+
	87	+
	88	+def _build_hf(, base: Any, args: Any, data_transform: Any, data_type: str, *_: Any) -> TokenizedDataset:
	89	+ """Shared loader for ``hf_datasets`` and ``json_file``.
	90	+
	91	+ Tokenises via ``data_transform`` when provided, then filters away empty
	92	+ sequences and wraps in :class:`TokenizedDataset`. Updates
	93	+ ``base.state.max_steps`` to ``min(cfg.max_steps, len/global_bs)`` so
	94	+ epoch boundaries stay consistent with the data on hand.
	95	+ """
	96	+ logger.info(
	97	+ "Loading dataset: type=%s, path=%s", data_type, args.data.train_path,
	98	+ )
	99	+ ds = _maybe_truncate(_load_raw(args, data_type), args)
	100	+
	101	+ if data_transform is not None:
	102	+ ds = ds.map(
	103	+ data_transform,
	104	+ batched=True,
	105	+ remove_columns=ds.column_names,
	106	+ desc="Tokenizing",
	107	+ )
	108	+ # Drop empty rows only when the dataset is already tokenized — without a
	109	+ # transform a raw text dataset has no ``input_ids`` column and the filter
	110	+ # predicate would raise ``KeyError`` instead of loading.
	111	+ if "input_ids" in ds.column_names:
	112	+ ds = ds.filter(lambda x: len(x["input_ids"]) > 0)
	113	+
	114	+ wrapped = TokenizedDataset(ds)
	115	+ # max_steps clipping must happen pre-dataloader so the train loop's
	116	+ # epoch count matches the data on hand; the trainer reads
	117	+ # ``base.state.max_steps`` further down the build chain. The cap is
	118	+ # the TOTAL step budget across all epochs — clipping to a single
	119	+ # epoch would silently truncate multi-epoch training.
	120	+ num_epochs = max(int(args.train.num_train_epochs or 1), 1)
	121	+ base.state.max_steps = min(
	122	+ args.train.max_steps,
	123	+ num_epochs * (len(wrapped) // max(args.train.global_batch_size, 1)),
	124	+ )
	125	+ logger.info(
	126	+ "Dataset ready: %d samples, max_steps=%d",
	127	+ len(wrapped), base.state.max_steps,
	128	+ )
	129	+ return wrapped
	130	+
	131	+
	132	+@DATASET_REGISTRY.register("hf_datasets")
	133	+def build_hf_datasets(**kwargs: Any) -> TokenizedDataset:
	134	+ """Build an HF hub / arrow dataset."""
	135	+ return _build_hf(data_type="hf_datasets", **kwargs)
	136	+
	137	+
	138	+@DATASET_REGISTRY.register("json_file")
	139	+def build_json_file(**kwargs: Any) -> TokenizedDataset:
	140	+ """Build a local Alpaca-style ``.json`` / ``.jsonl`` dataset."""
	141	+ return _build_hf(data_type="json_file", **kwargs)

Ahyper_parallel/data/megatron/__init__.py+48-0

@@ -0,0 +1,48 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Megatron ``.bin`` / ``.idx`` dataset support.
	16	+
	17	+Public API:
	18	+
	19	+- :class:`IndexedDataset` + :class:`IndexedDatasetBuilder` — wire-format
	20	+ compatible reader/writer for Megatron-LM's ``.bin`` / ``.idx``.
	21	+- :class:`GPTDataset` — fixed-length sampling on top of an IndexedDataset.
	22	+- :class:`BlendableDataset` — weighted multi-source blend.
	23	+
	24	+Importing this package registers ``data.type='megatron'`` with
	25	+:data:`hyper_parallel.data.DATASET_REGISTRY` via ``builder``.
	26	+"""
	27	+from hyper_parallel.data.megatron.blendable_dataset import BlendableDataset
	28	+from hyper_parallel.data.megatron.gpt_dataset import GPTDataset
	29	+from hyper_parallel.data.megatron.indexed_dataset import (
	30	+ IndexedDataset,
	31	+ IndexedDatasetBuilder,
	32	+ get_bin_path,
	33	+ get_idx_path,
	34	+ strip_suffix,
	35	+)
	36	+
	37	+from hyper_parallel.data.megatron import builder # noqa: F401 # registers builder
	38	+
	39	+
	40	+__all__ = [
	41	+ "BlendableDataset",
	42	+ "GPTDataset",
	43	+ "IndexedDataset",
	44	+ "IndexedDatasetBuilder",
	45	+ "get_bin_path",
	46	+ "get_idx_path",
	47	+ "strip_suffix",
	48	+]

Ahyper_parallel/data/megatron/blendable_dataset.py+138-0

@@ -0,0 +1,138 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Blend multiple datasets with per-source weights.
	16	+
	17	+Given ``N`` sub-datasets and ``N`` non-negative weights ``w_i``, a sample at
	18	+global index ``g`` is read from sub-dataset ``dataset_index[g]`` at local
	19	+index ``dataset_sample_index[g]``. Both indices are precomputed once (and
	20	+kept in memory) using a deterministic deficit scheduler, so the same
	21	+``(weights, num_samples)`` tuple always yields the same blend. ``seed`` is
	22	+accepted only for API compatibility.
	23	+
	24	+Use this to mix corpora — e.g. 30 % code, 70 % web — while keeping the
	25	+downstream :class:`GPTDataset` API intact (each sub-dataset can be a
	26	+GPTDataset itself, so blending stacks cleanly).
	27	+"""
	28	+import logging
	29	+from typing import Any, Dict, List, Sequence
	30	+
	31	+import numpy as np
	32	+from torch.utils.data import Dataset
	33	+
	34	+
	35	+logger = logging.getLogger(__name__)
	36	+
	37	+
	38	+def _build_blend_indices(
	39	+ dataset_sizes: Sequence[int],
	40	+ weights: Sequence[float],
	41	+ num_samples: int,
	42	+ seed: int, # pylint: disable=W0613
	43	+) -> tuple:
	44	+ """Return ``(dataset_index, dataset_sample_index)`` for a weighted blend.
	45	+
	46	+ At every global slot ``g``, pick the source whose realised fraction lags
	47	+ its target the most (i.e. ``target * (g+1) - realised`` is maximal).
	48	+ This guarantees the realised blend tracks ``weights`` exactly up to a
	49	+ one-sample rounding gap, with no run-to-run variance.
	50	+
	51	+ ``seed`` is accepted for API parity but unused; the algorithm is
	52	+ deterministic in ``(weights, num_samples)``.
	53	+ """
	54	+ if len(dataset_sizes) != len(weights):
	55	+ raise ValueError(
	56	+ f"dataset_sizes ({len(dataset_sizes)}) and weights ({len(weights)}) "
	57	+ f"must have the same length"
	58	+ )
	59	+ if any(w < 0 for w in weights):
	60	+ raise ValueError(f"weights must be non-negative, got {list(weights)}")
	61	+ total_w = float(sum(weights))
	62	+ if total_w <= 0:
	63	+ raise ValueError("weights must contain at least one non-zero value")
	64	+ norm_w = np.asarray([w / total_w for w in weights], dtype=np.float64)
	65	+ sizes = np.asarray([int(s) for s in dataset_sizes], dtype=np.int64)
	66	+ # A zero-length source with a non-zero weight would be selected by the
	67	+ # greatest-error rule and then hit ``counters[d] % sizes[d]`` → divide by
	68	+ # zero. Reject up-front (a source contributing samples must have samples).
	69	+ empty_weighted = [i for i, (s, w) in enumerate(zip(sizes, weights)) if s == 0 and w > 0]
	70	+ if empty_weighted:
	71	+ raise ValueError(
	72	+ f"BlendableDataset sources {empty_weighted} have a non-zero weight "
	73	+ f"but zero length; an empty source cannot contribute samples"
	74	+ )
	75	+
	76	+ dataset_index = np.zeros(num_samples, dtype=np.int32)
	77	+ dataset_sample_index = np.zeros(num_samples, dtype=np.int64)
	78	+ realised = np.zeros(len(weights), dtype=np.float64)
	79	+ counters = np.zeros(len(weights), dtype=np.int64)
	80	+ for g in range(num_samples):
	81	+ # error[d] = how far behind source d's realised fraction is —
	82	+ # picking the argmax keeps every source within 1 sample of its
	83	+ # target share at every prefix (Megatron's invariant).
	84	+ error = norm_w * (g + 1) - realised
	85	+ d = int(np.argmax(error))
	86	+ dataset_index[g] = d
	87	+ # Modulo wrap so a small sub-dataset reused inside a long blend
	88	+ # cycles through its samples in order instead of indexing OOB.
	89	+ dataset_sample_index[g] = counters[d] % sizes[d]
	90	+ realised[d] += 1.0
	91	+ counters[d] += 1
	92	+ return dataset_index, dataset_sample_index
	93	+
	94	+
	95	+class BlendableDataset(Dataset):
	96	+ """Weighted blend over a list of datasets.
	97	+
	98	+ Args:
	99	+ datasets: Concrete ``torch.utils.data.Dataset`` objects; each must
	100	+ implement ``__len__`` and ``__getitem__``. Typically these are
	101	+ :class:`GPTDataset` instances.
	102	+ weights: One non-negative weight per dataset (auto-normalised).
	103	+ num_samples: Total length exposed by the blend.
	104	+ seed: Accepted for API compatibility; index construction is
	105	+ deterministic and does not use random sampling.
	106	+
	107	+ Note:
	108	+ At every prefix, each source's realised count tracks its weighted
	109	+ target within the scheduler's one-sample rounding bound.
	110	+ """
	111	+
	112	+ def __init__(
	113	+ self,
	114	+ datasets: List[Dataset],
	115	+ weights: Sequence[float],
	116	+ num_samples: int,
	117	+ seed: int = 1234,
	118	+ ) -> None:
	119	+ if not datasets:
	120	+ raise ValueError("BlendableDataset requires at least one dataset")
	121	+ self.datasets = datasets
	122	+ sizes = [len(d) for d in datasets]
	123	+ self.num_samples = int(num_samples)
	124	+ self.dataset_index, self.dataset_sample_index = _build_blend_indices(
	125	+ sizes, weights, self.num_samples, seed,
	126	+ )
	127	+ logger.info(
	128	+ "BlendableDataset built: %d sub-datasets, weights=%s, total samples=%d",
	129	+ len(datasets), list(weights), self.num_samples,
	130	+ )
	131	+
	132	+ def __len__(self) -> int:
	133	+ return self.num_samples
	134	+
	135	+ def __getitem__(self, idx: int) -> Dict[str, Any]:
	136	+ d = int(self.dataset_index[idx])
	137	+ s = int(self.dataset_sample_index[idx])
	138	+ return self.datasets[d][s]

Ahyper_parallel/data/megatron/builder.py+165-0

@@ -0,0 +1,165 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Megatron dataset builder — wires config → IndexedDataset → GPTDataset (+ blend).
	16	+
	17	+Config surface (``args.data.*``):
	18	+
	19	+- ``train_path``: either a single path-prefix (no suffix) pointing to one
	20	+ ``.bin``/``.idx`` pair, OR a flat ``"w1 path1 w2 path2 ..."`` string /
	21	+ list for :class:`BlendableDataset`.
	22	+- ``max_seq_len``: per-sample token count (input plus implicit label shift).
	23	+- ``megatron_seed`` (optional, default ``train.seed``): RNG seed for the
	24	+ document / sample / shuffle indices.
	25	+- ``pad_token_id`` (optional, default 0): right-pad for the tail sample.
	26	+- ``eod_token_id`` (optional, default None) + ``eod_mask_loss`` (optional,
	27	+ default False): mask EOD tokens out of the loss.
	28	+
	29	+Sample count: ``num_samples = max_steps * global_batch_size`` so the
	30	+distributed sampler can drain the dataset without short cycles. Unlike the
	31	+HF builder, this builder does NOT modify ``base.state.max_steps``: Megatron
	32	+datasets are typically larger than ``num_samples``, so the configured
	33	+``max_steps`` is used as-is with no truncation.
	34	+"""
	35	+import logging
	36	+from typing import Any, List, Tuple
	37	+
	38	+from hyper_parallel.data.megatron.blendable_dataset import BlendableDataset
	39	+from hyper_parallel.data.megatron.gpt_dataset import GPTDataset
	40	+from hyper_parallel.data.megatron.indexed_dataset import IndexedDataset, strip_suffix
	41	+from hyper_parallel.data.registry import DATASET_REGISTRY
	42	+
	43	+
	44	+logger = logging.getLogger(__name__)
	45	+
	46	+
	47	+def _parse_blend(raw: Any) -> List[Tuple[float, str]]:
	48	+ """Parse a blend spec into ``[(weight, prefix), ...]``.
	49	+
	50	+ Supported shapes:
	51	+
	52	+ - ``str`` of whitespace-separated tokens: ``"0.3 /data/a 0.7 /data/b"``
	53	+ - ``list[float \| str]`` in the same alternating order
	54	+ - ``list[list]`` of ``[weight, prefix]`` pairs
	55	+ """
	56	+ if isinstance(raw, str):
	57	+ toks = raw.split()
	58	+ elif isinstance(raw, (list, tuple)):
	59	+ # Already in pair form?
	60	+ if raw and isinstance(raw[0], (list, tuple)):
	61	+ return [(float(w), str(p)) for w, p in raw]
	62	+ toks = list(raw)
	63	+ else:
	64	+ raise ValueError(f"Unsupported blend spec type: {type(raw)}")
	65	+
	66	+ if len(toks) % 2 != 0:
	67	+ raise ValueError(
	68	+ f"Blend spec must have an even number of tokens "
	69	+ f"(alternating weight + prefix); got {len(toks)}: {toks}"
	70	+ )
	71	+ out: List[Tuple[float, str]] = []
	72	+ for i in range(0, len(toks), 2):
	73	+ out.append((float(toks[i]), str(toks[i + 1])))
	74	+ return out
	75	+
	76	+
	77	+def _looks_like_blend(train_path: Any) -> bool:
	78	+ """Return ``True`` when ``train_path`` is a multi-source blend spec.
	79	+
	80	+ A blend spec always starts with a numeric weight (``"0.3 /a 0.7 /b"`` or
	81	+ a list of pairs); a single corpus is a bare path-prefix. Keying off "the
	82	+ first token parses as a float" (rather than "contains whitespace") lets a
	83	+ single path that legitimately contains spaces still be read as one source.
	84	+ """
	85	+ if isinstance(train_path, (list, tuple)):
	86	+ return True
	87	+ if not isinstance(train_path, str):
	88	+ return False
	89	+ toks = train_path.split()
	90	+ if len(toks) < 2:
	91	+ return False
	92	+ try:
	93	+ float(toks[0])
	94	+ except ValueError:
	95	+ return False
	96	+ return True
	97	+
	98	+
	99	+@DATASET_REGISTRY.register("megatron")
	100	+def build_megatron(, base: Any, args: Any, *_: Any) -> Any:
	101	+ """Build a Megatron ``.bin``/``.idx`` dataset (single source or blend)."""
	102	+ del base
	103	+ data_cfg = args.data
	104	+ train_cfg = args.train
	105	+ train_path = data_cfg.train_path
	106	+ if not train_path:
	107	+ raise ValueError("data.train_path is required when data.type='megatron'")
	108	+
	109	+ seq_length = int(data_cfg.max_seq_len)
	110	+ # ``megatron_seed`` may be present-but-None on a default DataConfig; fall
	111	+ # through to ``train.seed`` in that case so the typed config doesn't have
	112	+ # to invent a sentinel.
	113	+ megatron_seed = data_cfg.megatron_seed
	114	+ if megatron_seed is None:
	115	+ megatron_seed = train_cfg.seed
	116	+ seed = int(megatron_seed)
	117	+ pad_token_id = int(data_cfg.pad_token_id)
	118	+ eod_token_id = data_cfg.eod_token_id
	119	+ eod_mask_loss = bool(data_cfg.eod_mask_loss)
	120	+ mmap_bin = bool(data_cfg.mmap_bin_files)
	121	+
	122	+ num_samples = int(train_cfg.max_steps * train_cfg.global_batch_size)
	123	+ if num_samples <= 0:
	124	+ raise ValueError(
	125	+ f"num_samples = max_steps * global_batch_size must be > 0; "
	126	+ f"got max_steps={train_cfg.max_steps}, global_bs={train_cfg.global_batch_size}"
	127	+ )
	128	+
	129	+ # Single source: simplest path, no blend needed.
	130	+ if not _looks_like_blend(train_path):
	131	+ prefix = strip_suffix(train_path)
	132	+ idx_ds = IndexedDataset(prefix, mmap=mmap_bin)
	133	+ gpt_ds = GPTDataset(
	134	+ idx_ds, num_samples=num_samples, seq_length=seq_length,
	135	+ seed=seed, pad_token_id=pad_token_id,
	136	+ eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id,
	137	+ )
	138	+ logger.info(
	139	+ "Megatron dataset built: prefix=%s docs=%d sequences=%d samples=%d seq_len=%d",
	140	+ prefix, idx_ds.num_documents, len(idx_ds), len(gpt_ds), seq_length,
	141	+ )
	142	+ return gpt_ds
	143	+
	144	+ # Blend: parse, build per-source GPTDataset, wrap in BlendableDataset.
	145	+ pairs = _parse_blend(train_path)
	146	+ weights = [w for w, _ in pairs]
	147	+ sub_datasets = []
	148	+ for w, p in pairs:
	149	+ prefix = strip_suffix(p)
	150	+ idx_ds = IndexedDataset(prefix, mmap=mmap_bin)
	151	+ # Each sub-dataset is sized to cover its share of the blend; we ask
	152	+ # for the full ``num_samples`` here because blending picks indices
	153	+ # modulo the sub-dataset length.
	154	+ sub_datasets.append(
	155	+ GPTDataset(
	156	+ idx_ds, num_samples=num_samples, seq_length=seq_length,
	157	+ seed=seed, pad_token_id=pad_token_id,
	158	+ eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id,
	159	+ )
	160	+ )
	161	+ logger.info(
	162	+ " blend source w=%g prefix=%s docs=%d sequences=%d",
	163	+ w, prefix, idx_ds.num_documents, len(idx_ds),
	164	+ )
	165	+ return BlendableDataset(sub_datasets, weights, num_samples=num_samples, seed=seed)

Ahyper_parallel/data/megatron/gpt_dataset.py+331-0

@@ -0,0 +1,331 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""``GPTDataset`` — fixed-length samples drawn from a Megatron IndexedDataset.
	16	+
	17	+Behaviourally aligned with Megatron-LM's ``GPTDataset``:
	18	+
	19	+1. ``document_index`` — a shuffled flat array repeating every document for
	20	+ ``num_epochs`` epochs. Documents stay whole; only their visit order is
	21	+ randomised.
	22	+2. ``sample_index`` — for sample ``i`` the entry ``[doc_id, doc_offset]``
	23	+ points to the document and within-document token offset where that
	24	+ sample starts. Built by sequentially walking the document stream and
	25	+ slicing every ``seq_length`` tokens.
	26	+3. ``shuffle_index`` — a random permutation applied on top, so sample
	27	+ ``__getitem__(i)`` actually reads sample ``shuffle_index[i]``.
	28	+
	29	+A sample may span multiple documents (Megatron's concatenated-doc style);
	30	+in that case we glue the relevant slices together until we reach
	31	+``seq_length + 1`` tokens (one extra so labels can be shifted, exactly
	32	+Megatron's ``add_extra_token_to_sequence=1``). If the very last sample
	33	+runs out of documents, the tail is right-padded with ``pad_token_id``.
	34	+
	35	+Pure Python+numpy. No Cython helpers — the index construction runs once
	36	+per ``(seed, prefix, num_samples, seq_length, num_epochs)`` tuple and is
	37	+cached to disk under ``<prefix>_cache/`` so subsequent runs reuse it.
	38	+"""
	39	+import functools
	40	+import hashlib
	41	+import logging
	42	+import os
	43	+import tempfile
	44	+from typing import Any, Dict, Optional, Tuple
	45	+
	46	+import numpy as np
	47	+import torch
	48	+from torch.utils.data import Dataset
	49	+
	50	+from hyper_parallel.data.megatron.indexed_dataset import IndexedDataset
	51	+
	52	+
	53	+logger = logging.getLogger(__name__)
	54	+
	55	+
	56	+def _build_document_index(num_documents: int, num_epochs: int, rng: np.random.RandomState) -> np.ndarray:
	57	+ """Return a globally-shuffled, ``num_epochs``-repeated array of document indices.
	58	+
	59	+ Matches Megatron-LM's ``separate_last_epoch=False`` default: tile the
	60	+ document range across ``num_epochs`` and run a single global shuffle.
	61	+ This intentionally lets a document appear several times before another
	62	+ appears once — that randomness is part of the canonical training
	63	+ distribution and the realised epoch-boundary noise is what Megatron's
	64	+ schedulers expect.
	65	+ """
	66	+ doc_idx = np.tile(np.arange(num_documents, dtype=np.int32), num_epochs)
	67	+ rng.shuffle(doc_idx)
	68	+ return doc_idx
	69	+
	70	+
	71	+def _build_sample_index(
	72	+ sizes: np.ndarray,
	73	+ doc_idx: np.ndarray,
	74	+ seq_length: int,
	75	+ tokens_per_epoch: int,
	76	+ num_epochs: int,
	77	+) -> np.ndarray:
	78	+ """Construct the ``[num_samples + 1, 2]`` (doc_id, doc_offset) sample table.
	79	+
	80	+ Walks the per-document token stream in ``doc_idx`` order, marking the
	81	+ start of every ``seq_length``-window. Each sample includes one extra
	82	+ token so labels can be shifted by 1; the helper keeps the
	83	+ pre-Megatron-Cython slow-path semantics (no boundary correction —
	84	+ multi-document samples are reassembled by :meth:`GPTDataset._read_sample`).
	85	+ """
	86	+ tokens_to_consume = num_epochs * tokens_per_epoch - 1 # -1 to keep ``+1`` extra token in reach
	87	+ sample_stride = seq_length
	88	+ num_samples = tokens_to_consume // sample_stride
	89	+ # int64: column 0 indexes ``doc_idx`` whose length is
	90	+ # ``num_documents * num_epochs`` — int32 would wrap past ~2.1e9 slots on
	91	+ # large corpora trained for many epochs, silently selecting the wrong
	92	+ # document.
	93	+ sample_idx = np.zeros((num_samples + 1, 2), dtype=np.int64)
	94	+
	95	+ cur_doc_index = 0
	96	+ cur_doc_offset = 0
	97	+ cur_doc_remaining = int(sizes[doc_idx[cur_doc_index]])
	98	+ sample_idx[0, 0] = cur_doc_index
	99	+ sample_idx[0, 1] = cur_doc_offset
	100	+ for sample_id in range(1, num_samples + 1):
	101	+ remaining = sample_stride
	102	+ while remaining > 0:
	103	+ if cur_doc_remaining > remaining:
	104	+ cur_doc_offset += remaining
	105	+ cur_doc_remaining -= remaining
	106	+ remaining = 0
	107	+ else:
	108	+ remaining -= cur_doc_remaining
	109	+ cur_doc_index += 1
	110	+ if cur_doc_index >= doc_idx.shape[0]:
	111	+ # Document stream exhausted — record where we landed
	112	+ # and stop. Caller will pad the tail.
	113	+ cur_doc_remaining = 0
	114	+ break
	115	+ cur_doc_offset = 0
	116	+ cur_doc_remaining = int(sizes[doc_idx[cur_doc_index]])
	117	+ sample_idx[sample_id, 0] = cur_doc_index
	118	+ sample_idx[sample_id, 1] = cur_doc_offset
	119	+ return sample_idx
	120	+
	121	+
	122	+def _build_shuffle_index(num_samples: int, rng: np.random.RandomState) -> np.ndarray:
	123	+ """Random permutation of ``[0, num_samples)``."""
	124	+ shuffle = np.arange(num_samples, dtype=np.int64)
	125	+ rng.shuffle(shuffle)
	126	+ return shuffle
	127	+
	128	+
	129	+def _layout_fingerprint(doc_sizes: np.ndarray) -> str:
	130	+ """Hash the per-document token-count array.
	131	+
	132	+ ``doc_sizes`` is exactly what determines every ``sample_idx`` offset, so a
	133	+ corpus regenerated at the same path with a different document layout — even
	134	+ one preserving document count and total token count — produces a different
	135	+ fingerprint and invalidates the stale cache.
	136	+ """
	137	+ return hashlib.sha1(np.ascontiguousarray(doc_sizes, dtype=np.int64).tobytes()).hexdigest()[:16]
	138	+
	139	+
	140	+def _cache_key(
	141	+ prefix: str,
	142	+ seq_length: int,
	143	+ num_samples: int,
	144	+ num_epochs: int,
	145	+ seed: int,
	146	+ layout_fp: str,
	147	+) -> str:
	148	+ """Stable filename suffix for the ``(doc, sample, shuffle)`` index trio.
	149	+
	150	+ ``layout_fp`` fingerprints the underlying corpus layout (see
	151	+ :func:`_layout_fingerprint`) so that regenerating the ``.bin``/``.idx``
	152	+ under the same path invalidates a stale cache instead of silently
	153	+ replaying indices computed against the old token stream.
	154	+ """
	155	+ h = hashlib.sha1(
	156	+ f"{prefix}\|{seq_length}\|{num_samples}\|{num_epochs}\|{seed}\|{layout_fp}".encode("utf-8"),
	157	+ ).hexdigest()[:16]
	158	+ return h
	159	+
	160	+
	161	+class GPTDataset(Dataset):
	162	+ """Megatron-style GPT dataset on top of an :class:`IndexedDataset`.
	163	+
	164	+ Args:
	165	+ indexed_dataset: Underlying token-stream reader.
	166	+ num_samples: How many fixed-length samples to expose. Typically
	167	+ ``max_steps * global_batch_size`` so the dataloader never
	168	+ runs dry mid-training.
	169	+ seq_length: Tokens per ``input_ids`` (excluding the extra label token).
	170	+ seed: RNG seed for the document / shuffle permutations.
	171	+ pad_token_id: Token used to right-pad the very last sample when
	172	+ the document stream is exhausted before ``seq_length + 1``
	173	+ tokens are available.
	174	+ eod_mask_loss: When ``True``, the EOD token's label is replaced by
	175	+ ``-100`` so it does not contribute to the loss.
	176	+ eod_token_id: EOD token id; ignored when ``eod_mask_loss`` is False.
	177	+ cache_dir: Where to cache the (doc, sample, shuffle) index arrays.
	178	+ Defaults to ``<indexed_dataset.path_prefix>_cache``.
	179	+
	180	+ Returns from ``__getitem__``:
	181	+ ``{"input_ids": Tensor[seq_length], "labels": Tensor[seq_length]}``
	182	+ — labels is a copy of input_ids (the model handles next-token
	183	+ shifting internally, matching the rest of the trainer's
	184	+ dataset contract).
	185	+ """
	186	+
	187	+ def __init__(
	188	+ self,
	189	+ indexed_dataset: IndexedDataset,
	190	+ num_samples: int,
	191	+ seq_length: int,
	192	+ seed: int = 1234,
	193	+ *,
	194	+ pad_token_id: int = 0,
	195	+ eod_mask_loss: bool = False,
	196	+ eod_token_id: Optional[int] = None,
	197	+ cache_dir: Optional[str] = None,
	198	+ ) -> None:
	199	+ if num_samples <= 0:
	200	+ raise ValueError(f"num_samples must be > 0, got {num_samples}")
	201	+ if seq_length <= 0:
	202	+ raise ValueError(f"seq_length must be > 0, got {seq_length}")
	203	+
	204	+ self.indexed_dataset = indexed_dataset
	205	+ self.num_samples = int(num_samples)
	206	+ self.seq_length = int(seq_length)
	207	+ self.seed = int(seed)
	208	+ self.pad_token_id = int(pad_token_id)
	209	+ self.eod_mask_loss = bool(eod_mask_loss)
	210	+ self.eod_token_id = eod_token_id
	211	+
	212	+ self.document_index, self.sample_index, self.shuffle_index = self._build_indices(cache_dir)
	213	+ # Adjacent samples often touch the same document — cache the concat
	214	+ # result for the last few requested doc ids so straddling reads
	215	+ # don't re-walk the IndexedDataset.
	216	+ self._cached_doc = functools.lru_cache(maxsize=128)(self._read_doc)
	217	+
	218	+ def _build_indices(self, cache_dir: Optional[str]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
	219	+ """Build (or load from cache) the three Megatron-style indices."""
	220	+ sizes = self.indexed_dataset.sequence_lengths
	221	+ num_docs = self.indexed_dataset.num_documents
	222	+ if num_docs == 0:
	223	+ raise ValueError("IndexedDataset has no documents to sample from")
	224	+ # Build a per-document length array using the existing
	225	+ # ``document_indices`` boundaries.
	226	+ doc_boundaries = self.indexed_dataset.document_indices
	227	+ doc_sizes = np.zeros(num_docs, dtype=np.int64)
	228	+ for d in range(num_docs):
	229	+ start = int(doc_boundaries[d])
	230	+ end = int(doc_boundaries[d + 1])
	231	+ doc_sizes[d] = int(sizes[start:end].sum())
	232	+
	233	+ tokens_per_epoch = int(doc_sizes.sum())
	234	+ if tokens_per_epoch == 0:
	235	+ raise ValueError("IndexedDataset has zero total tokens")
	236	+ # ``+1`` because each sample reads ``seq_length + 1`` tokens (input
	237	+ # plus the next-token label). Ceil so we always cover ``num_samples``.
	238	+ tokens_required = self.num_samples * self.seq_length + 1
	239	+ num_epochs = max(1, int(np.ceil(tokens_required / tokens_per_epoch)))
	240	+
	241	+ cache_dir = self._resolve_cache_dir(cache_dir)
	242	+ key = _cache_key(
	243	+ self.indexed_dataset.path_prefix, self.seq_length,
	244	+ self.num_samples, num_epochs, self.seed,
	245	+ _layout_fingerprint(doc_sizes),
	246	+ )
	247	+ doc_path = os.path.join(cache_dir, f"doc_idx_{key}.npy")
	248	+ sample_path = os.path.join(cache_dir, f"sample_idx_{key}.npy")
	249	+ shuffle_path = os.path.join(cache_dir, f"shuffle_idx_{key}.npy")
	250	+
	251	+ if os.path.isfile(doc_path) and os.path.isfile(sample_path) and os.path.isfile(shuffle_path):
	252	+ return np.load(doc_path), np.load(sample_path), np.load(shuffle_path)
	253	+
	254	+ rng = np.random.RandomState(self.seed)
	255	+ doc_idx = _build_document_index(num_docs, num_epochs, rng)
	256	+ sample_idx = _build_sample_index(
	257	+ doc_sizes, doc_idx, self.seq_length, tokens_per_epoch, num_epochs,
	258	+ )
	259	+ # ``sample_idx`` has at most ``num_samples + 1`` rows; clip in case
	260	+ # we built slightly more (rounded up by tokens_per_epoch / seq_len).
	261	+ usable_samples = max(sample_idx.shape[0] - 1, 1)
	262	+ shuffle_idx = _build_shuffle_index(min(self.num_samples, usable_samples), rng)
	263	+ np.save(doc_path, doc_idx)
	264	+ np.save(sample_path, sample_idx)
	265	+ np.save(shuffle_path, shuffle_idx)
	266	+ return doc_idx, sample_idx, shuffle_idx
	267	+
	268	+ def _resolve_cache_dir(self, user_cache_dir: Optional[str]) -> str:
	269	+ """Return a writable cache directory, falling back to ``$TMPDIR``.
	270	+
	271	+ Defaulting to ``<prefix>_cache`` next to the corpus is convenient,
	272	+ but production datasets often sit on read-only shared mounts where
	273	+ ``os.makedirs`` would raise ``EROFS``. Detect and silently relocate.
	274	+ """
	275	+ candidate = user_cache_dir or (self.indexed_dataset.path_prefix + "_cache")
	276	+ try:
	277	+ os.makedirs(candidate, exist_ok=True)
	278	+ # Probe writability — ``os.makedirs`` succeeds on directories
	279	+ # that exist read-only.
	280	+ probe = os.path.join(candidate, ".hp_cache_probe")
	281	+ with open(probe, "w", encoding="utf-8") as f:
	282	+ f.write("")
	283	+ os.remove(probe)
	284	+ return candidate
	285	+ except OSError as exc:
	286	+ fallback = os.path.join(tempfile.gettempdir(), "hp_megatron_cache")
	287	+ os.makedirs(fallback, exist_ok=True)
	288	+ logger.warning(
	289	+ "GPTDataset cache dir %s not writable (%s); using fallback %s",
	290	+ candidate, exc, fallback,
	291	+ )
	292	+ return fallback
	293	+
	294	+ def __len__(self) -> int:
	295	+ return min(self.num_samples, int(self.shuffle_index.size))
	296	+
	297	+ def _read_sample(self, sample_pos: int) -> np.ndarray:
	298	+ """Concatenate the ``seq_length + 1`` tokens for sample ``sample_pos``."""
	299	+ start_doc_idx, start_offset = self.sample_index[sample_pos]
	300	+ end_doc_idx, end_offset = self.sample_index[sample_pos + 1]
	301	+ parts = []
	302	+ if start_doc_idx == end_doc_idx:
	303	+ doc_tokens = self._cached_doc(int(self.document_index[start_doc_idx]))
	304	+ parts.append(doc_tokens[start_offset:end_offset + 1])
	305	+ else:
	306	+ doc_tokens = self._cached_doc(int(self.document_index[start_doc_idx]))
	307	+ parts.append(doc_tokens[start_offset:])
	308	+ for doc_id in range(start_doc_idx + 1, end_doc_idx):
	309	+ parts.append(self._cached_doc(int(self.document_index[doc_id])))
	310	+ tail_doc_tokens = self._cached_doc(int(self.document_index[end_doc_idx]))
	311	+ parts.append(tail_doc_tokens[:end_offset + 1])
	312	+ out = np.concatenate(parts) if len(parts) > 1 else parts[0]
	313	+ needed = self.seq_length + 1
	314	+ if out.size >= needed:
	315	+ return out[:needed]
	316	+ # Pad short tail (only happens for the very last sample).
	317	+ pad = np.full(needed - out.size, self.pad_token_id, dtype=out.dtype)
	318	+ return np.concatenate([out, pad])
	319	+
	320	+ def _read_doc(self, doc_id: int) -> np.ndarray:
	321	+ """All tokens of ``doc_id`` as a single 1-D array (memoised)."""
	322	+ return self.indexed_dataset.get_document(doc_id)
	323	+
	324	+ def __getitem__(self, idx: int) -> Dict[str, Any]:
	325	+ actual = int(self.shuffle_index[idx])
	326	+ tokens = self._read_sample(actual)
	327	+ input_ids = torch.from_numpy(tokens[: self.seq_length].astype(np.int64))
	328	+ labels = input_ids.clone()
	329	+ if self.eod_mask_loss and self.eod_token_id is not None:
	330	+ labels[input_ids == int(self.eod_token_id)] = -100
	331	+ return {"input_ids": input_ids, "labels": labels}

Ahyper_parallel/data/megatron/indexed_dataset.py+364-0

@@ -0,0 +1,364 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Megatron-LM ``.bin`` / ``.idx`` reader and writer.
	16	+
	17	+Wire format (matches Megatron-LM ``mmap_indexed_dataset.py``):
	18	+
	19	+``<prefix>.idx``:
	20	+- 9-byte magic ``b"MMIDIDX\\x00\\x00"``
	21	+- ``uint64`` version (=1)
	22	+- ``uint8`` dtype code (1=uint8, 2=int8, 3=int16, 4=int32, 5=int64,
	23	+ 6=float64, 7=float32, 8=uint16)
	24	+- ``uint64`` sequence count ``N``
	25	+- ``uint64`` document count ``D``
	26	+- ``int32 [N]`` sequence lengths (in tokens)
	27	+- ``int64 [N]`` sequence pointers (byte offset into ``.bin``)
	28	+- ``int64 [D + 1]`` document indices (sequence-index boundaries; first is 0,
	29	+ last is ``N``)
	30	+- Optional ``int8 [N]`` sequence modes for multimodal datasets — ignored
	31	+ here (LM-only path)
	32	+
	33	+``<prefix>.bin``: raw concatenation of all sequences in the declared dtype.
	34	+
	35	+We expose:
	36	+
	37	+- :func:`get_bin_path` / :func:`get_idx_path` — canonical suffixing logic.
	38	+- :class:`IndexedDataset` — memory-mapped ``.bin`` payload with in-memory
	39	+ ``.idx`` metadata; ``__getitem__`` returns ``np.ndarray``.
	40	+- :class:`IndexedDatasetBuilder` — sequential writer for tests, fixtures,
	41	+ and ad-hoc preprocessing.
	42	+
	43	+Pure ``numpy`` — no Cython helpers — so this works wherever ``numpy`` runs.
	44	+The data path stays zero-copy: token slices view the memory-mapped ``.bin``;
	45	+the smaller ``.idx`` arrays are read into RAM once at open time.
	46	+"""
	47	+import logging
	48	+import os
	49	+import struct
	50	+from typing import Dict, List, Optional, Tuple
	51	+
	52	+import numpy as np
	53	+
	54	+
	55	+logger = logging.getLogger(__name__)
	56	+
	57	+
	58	+# Megatron magic header — kept byte-identical so a ``.idx`` produced by
	59	+# Megatron-LM ``preprocess_data.py`` loads transparently here.
	60	+_MAGIC = b"MMIDIDX\x00\x00"
	61	+_VERSION = 1
	62	+
	63	+_DTYPES: Dict[int, np.dtype] = {
	64	+ 1: np.dtype(np.uint8),
	65	+ 2: np.dtype(np.int8),
	66	+ 3: np.dtype(np.int16),
	67	+ 4: np.dtype(np.int32),
	68	+ 5: np.dtype(np.int64),
	69	+ 6: np.dtype(np.float64),
	70	+ 7: np.dtype(np.float32),
	71	+ 8: np.dtype(np.uint16),
	72	+}
	73	+_DTYPE_CODES: Dict[np.dtype, int] = {v: k for k, v in _DTYPES.items()}
	74	+
	75	+
	76	+def strip_suffix(prefix: str) -> str:
	77	+ """Return ``prefix`` with a trailing ``.bin``/``.idx`` removed (idempotent).
	78	+
	79	+ The single source of truth for path-prefix normalisation — :func:`get_bin_path`,
	80	+ :func:`get_idx_path` and the dataset builder all route through here so the
	81	+ suffix rule never diverges across call sites.
	82	+ """
	83	+ return prefix[:-4] if prefix.endswith((".bin", ".idx")) else prefix
	84	+
	85	+
	86	+def get_bin_path(prefix: str) -> str:
	87	+ """Return ``<prefix>.bin`` (strips an already-present ``.bin``/``.idx``)."""
	88	+ return strip_suffix(prefix) + ".bin"
	89	+
	90	+
	91	+def get_idx_path(prefix: str) -> str:
	92	+ """Return ``<prefix>.idx``."""
	93	+ return strip_suffix(prefix) + ".idx"
	94	+
	95	+
	96	+def dtype_code(dtype: np.dtype) -> int:
	97	+ """Return the on-disk dtype code for ``dtype``.
	98	+
	99	+ Raises:
	100	+ ValueError: When ``dtype`` is not part of the Megatron dtype table.
	101	+ """
	102	+ dtype = np.dtype(dtype)
	103	+ if dtype not in _DTYPE_CODES:
	104	+ raise ValueError(
	105	+ f"Unsupported Megatron indexed-dataset dtype: {dtype}. "
	106	+ f"Allowed: {sorted(d.name for d in _DTYPE_CODES)}"
	107	+ )
	108	+ return _DTYPE_CODES[dtype]
	109	+
	110	+
	111	+def code_dtype(code: int) -> np.dtype:
	112	+ """Inverse of :func:`dtype_code`."""
	113	+ try:
	114	+ return _DTYPES[code]
	115	+ except KeyError as exc:
	116	+ raise ValueError(f"Unknown Megatron dtype code: {code}") from exc
	117	+
	118	+
	119	+class _IndexReader:
	120	+ """In-memory metadata parsed from the ``.idx`` companion file.
	121	+
	122	+ Holds the four arrays Megatron-LM stores: lengths, pointers,
	123	+ document indices and optional modes. The index payload is read into
	124	+ compact numpy arrays; the much larger ``.bin`` remains memory-mapped by
	125	+ :class:`IndexedDataset`.
	126	+ """
	127	+
	128	+ def __init__(self, idx_path: str, *, multimodal: bool = False) -> None:
	129	+ if not os.path.isfile(idx_path):
	130	+ raise FileNotFoundError(f"Megatron index file not found: {idx_path}")
	131	+ with open(idx_path, "rb") as f:
	132	+ magic = f.read(len(_MAGIC))
	133	+ if magic != _MAGIC:
	134	+ raise ValueError(
	135	+ f"Not a Megatron indexed dataset (magic mismatch): "
	136	+ f"{idx_path}. Expected {_MAGIC!r}, got {magic!r}"
	137	+ )
	138	+ version = struct.unpack("<Q", f.read(8))[0]
	139	+ if version != _VERSION:
	140	+ raise ValueError(
	141	+ f"Unsupported indexed-dataset version: {version} "
	142	+ f"(only v{_VERSION} is recognised)"
	143	+ )
	144	+ code = struct.unpack("<B", f.read(1))[0]
	145	+ self.dtype = code_dtype(code)
	146	+ self.sequence_count = struct.unpack("<Q", f.read(8))[0]
	147	+ self.document_count = struct.unpack("<Q", f.read(8))[0]
	148	+
	149	+ # The index payload is small (4-12 bytes per sequence) — read it
	150	+ # into RAM rather than memmap. Avoids the fragility of slicing a
	151	+ # memmap with offsetted ``np.frombuffer`` while still being fast.
	152	+ self.sequence_lengths = np.frombuffer(
	153	+ f.read(self.sequence_count * 4), dtype=np.int32,
	154	+ )
	155	+ self.sequence_pointers = np.frombuffer(
	156	+ f.read(self.sequence_count * 8), dtype=np.int64,
	157	+ )
	158	+ self.document_indices = np.frombuffer(
	159	+ f.read(self.document_count * 8), dtype=np.int64,
	160	+ )
	161	+ if multimodal:
	162	+ self.sequence_modes = np.frombuffer(
	163	+ f.read(self.sequence_count), dtype=np.int8,
	164	+ )
	165	+ else:
	166	+ self.sequence_modes = None
	167	+
	168	+
	169	+class IndexedDataset:
	170	+ """Read-only Megatron ``.bin``/``.idx`` dataset.
	171	+
	172	+ Args:
	173	+ path_prefix: Path without the ``.bin``/``.idx`` suffix. Either
	174	+ suffix is also accepted and stripped.
	175	+ mmap: When ``True`` (default), the ``.bin`` is memory-mapped; reads
	176	+ return zero-copy views. Set to ``False`` to load the entire
	177	+ ``.bin`` into RAM — useful only for sub-MB fixtures in tests.
	178	+ multimodal: Whether to parse the optional ``sequence_modes`` array
	179	+ in the ``.idx`` (LM training leaves this off).
	180	+
	181	+ Example:
	182	+ >>> ds = IndexedDataset("/data/wiki") # noqa: SKIP
	183	+ >>> tokens = ds[0] # noqa: SKIP
	184	+ >>> doc_0 = ds.get_document(0) # noqa: SKIP
	185	+
	186	+ Note:
	187	+ Equivalent to Megatron-LM's ``MMapIndexedDataset`` but rewritten
	188	+ for the HyperParallel coding conventions; the on-disk format is
	189	+ byte-compatible.
	190	+ """
	191	+
	192	+ def __init__(self, path_prefix: str, *, mmap: bool = True, multimodal: bool = False) -> None:
	193	+ idx_path = get_idx_path(path_prefix)
	194	+ bin_path = get_bin_path(path_prefix)
	195	+ if not os.path.isfile(bin_path):
	196	+ raise FileNotFoundError(f"Megatron .bin not found: {bin_path}")
	197	+ self._index = _IndexReader(idx_path, multimodal=multimodal)
	198	+ if mmap:
	199	+ self._bin_buffer = np.memmap(bin_path, dtype=np.uint8, mode="r", order="C")
	200	+ else:
	201	+ with open(bin_path, "rb") as f:
	202	+ data = f.read()
	203	+ self._bin_buffer = np.frombuffer(data, dtype=np.uint8)
	204	+ self.path_prefix = path_prefix
	205	+
	206	+ # ------------------------------------------------------------------
	207	+ # Sequence-level access
	208	+ # ------------------------------------------------------------------
	209	+
	210	+ def __len__(self) -> int:
	211	+ return int(self._index.sequence_count)
	212	+
	213	+ @property
	214	+ def dtype(self) -> np.dtype:
	215	+ """Token dtype (numpy)."""
	216	+ return self._index.dtype
	217	+
	218	+ @property
	219	+ def sequence_lengths(self) -> np.ndarray:
	220	+ """int32[N] token-counts for each sequence."""
	221	+ return self._index.sequence_lengths
	222	+
	223	+ @property
	224	+ def document_indices(self) -> np.ndarray:
	225	+ """int64[D+1] sequence-index boundaries marking documents."""
	226	+ return self._index.document_indices
	227	+
	228	+ def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> np.ndarray:
	229	+ """Read a slice of sequence ``idx``.
	230	+
	231	+ Args:
	232	+ idx: Sequence index in ``[0, len)``.
	233	+ offset: Token offset within the sequence.
	234	+ length: Tokens to read; ``None`` reads to the end.
	235	+
	236	+ Returns:
	237	+ A ``np.ndarray`` of the configured dtype. The returned array is
	238	+ a view into the mmap when ``mmap=True``; callers must not write
	239	+ to it.
	240	+ """
	241	+ seq_len = int(self._index.sequence_lengths[idx])
	242	+ if length is None:
	243	+ length = seq_len - offset
	244	+ if offset < 0 or length < 0 or offset + length > seq_len:
	245	+ raise ValueError(
	246	+ f"Slice [{offset}:{offset + length}] out of bounds for "
	247	+ f"sequence {idx} of length {seq_len}"
	248	+ )
	249	+ item_size = self._index.dtype.itemsize
	250	+ byte_ptr = int(self._index.sequence_pointers[idx]) + offset * item_size
	251	+ # ``sequence_pointers`` are accumulated ``nbytes`` of items, so
	252	+ # ``byte_ptr`` is always a multiple of ``item_size`` and a typed
	253	+ # view is alignment-safe; the slice is a zero-copy view into the
	254	+ # mmap (or a plain ndarray slice in the RAM path).
	255	+ return self._bin_buffer[byte_ptr: byte_ptr + length * item_size].view(self._index.dtype)
	256	+
	257	+ def __getitem__(self, idx: int) -> np.ndarray:
	258	+ return self.get(idx)
	259	+
	260	+ # ------------------------------------------------------------------
	261	+ # Document-level access
	262	+ # ------------------------------------------------------------------
	263	+
	264	+ @property
	265	+ def num_documents(self) -> int:
	266	+ """Document count (``len(document_indices) - 1``).
	267	+
	268	+ On disk ``document_count`` is ``D + 1`` because Megatron stores the
	269	+ full boundary array including the trailing ``N`` terminator.
	270	+ """
	271	+ return int(self._index.document_count) - 1
	272	+
	273	+ def get_document(self, doc_idx: int) -> np.ndarray:
	274	+ """Concatenate all sequences belonging to ``doc_idx``."""
	275	+ if doc_idx < 0 or doc_idx >= self.num_documents:
	276	+ raise IndexError(
	277	+ f"document index {doc_idx} out of range [0, {self.num_documents})"
	278	+ )
	279	+ start = int(self._index.document_indices[doc_idx])
	280	+ end = int(self._index.document_indices[doc_idx + 1])
	281	+ parts = [self.get(i) for i in range(start, end)]
	282	+ if not parts:
	283	+ return np.array([], dtype=self._index.dtype)
	284	+ return np.concatenate(parts)
	285	+
	286	+
	287	+class IndexedDatasetBuilder:
	288	+ """Write a ``.bin``/``.idx`` pair sequentially.
	289	+
	290	+ Usage:
	291	+ builder = IndexedDatasetBuilder("/path/to/wiki", dtype=np.int32)
	292	+ builder.add_item(np.array([1, 2, 3, 4], dtype=np.int32))
	293	+ builder.end_document()
	294	+ builder.add_item(np.array([5, 6, 7], dtype=np.int32))
	295	+ builder.end_document()
	296	+ builder.finalize()
	297	+
	298	+ Args:
	299	+ path_prefix: Output path without suffix.
	300	+ dtype: Token dtype; must be in the Megatron dtype table.
	301	+
	302	+ Note:
	303	+ Designed for fixture generation and small-scale preprocessing.
	304	+ Streams tokens straight to the ``.bin`` and keeps lightweight
	305	+ Python lists for the index metadata until :meth:`finalize`.
	306	+ """
	307	+
	308	+ def __init__(self, path_prefix: str, *, dtype: np.dtype = np.int32) -> None:
	309	+ self.path_prefix = path_prefix
	310	+ self.dtype = np.dtype(dtype)
	311	+ # Make sure the dtype is supported before opening any file.
	312	+ dtype_code(self.dtype)
	313	+ self._bin_path = get_bin_path(path_prefix)
	314	+ self._idx_path = get_idx_path(path_prefix)
	315	+ os.makedirs(os.path.dirname(self._bin_path) or ".", exist_ok=True)
	316	+ self._bin_file = open(self._bin_path, "wb") # pylint: disable=R1732
	317	+ self._sequence_lengths: List[int] = []
	318	+ self._sequence_pointers: List[int] = []
	319	+ self._document_indices: List[int] = [0]
	320	+ self._byte_offset = 0
	321	+
	322	+ def add_item(self, tokens: np.ndarray) -> None:
	323	+ """Append one sequence of tokens to the dataset."""
	324	+ tokens = np.asarray(tokens, dtype=self.dtype)
	325	+ self._bin_file.write(tokens.tobytes(order="C"))
	326	+ self._sequence_lengths.append(int(tokens.size))
	327	+ self._sequence_pointers.append(self._byte_offset)
	328	+ self._byte_offset += int(tokens.nbytes)
	329	+
	330	+ def end_document(self) -> None:
	331	+ """Mark the current sequence boundary as a document boundary."""
	332	+ self._document_indices.append(len(self._sequence_lengths))
	333	+
	334	+ def finalize(self) -> Tuple[str, str]:
	335	+ """Close ``.bin`` and write the companion ``.idx``.
	336	+
	337	+ Returns:
	338	+ ``(idx_path, bin_path)`` for caller convenience.
	339	+
	340	+ Raises:
	341	+ ValueError: When no document boundary was emitted.
	342	+ """
	343	+ self._bin_file.close()
	344	+ # ``end_document`` must have been called at least once OR the writer
	345	+ # is empty; an empty dataset still needs a 1-element boundary list.
	346	+ if self._document_indices[-1] != len(self._sequence_lengths):
	347	+ self.end_document()
	348	+ seq_lengths = np.asarray(self._sequence_lengths, dtype=np.int32)
	349	+ seq_pointers = np.asarray(self._sequence_pointers, dtype=np.int64)
	350	+ doc_indices = np.asarray(self._document_indices, dtype=np.int64)
	351	+ with open(self._idx_path, "wb") as f:
	352	+ f.write(_MAGIC)
	353	+ f.write(struct.pack("<Q", _VERSION))
	354	+ f.write(struct.pack("<B", dtype_code(self.dtype)))
	355	+ f.write(struct.pack("<Q", seq_lengths.size))
	356	+ f.write(struct.pack("<Q", doc_indices.size))
	357	+ f.write(seq_lengths.tobytes())
	358	+ f.write(seq_pointers.tobytes())
	359	+ f.write(doc_indices.tobytes())
	360	+ logger.info(
	361	+ "IndexedDatasetBuilder finalized: %d sequences, %d documents -> %s + %s",
	362	+ seq_lengths.size, doc_indices.size, self._bin_path, self._idx_path,
	363	+ )
	364	+ return self._idx_path, self._bin_path

Ahyper_parallel/data/preset_pt.py+166-0

@@ -0,0 +1,166 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""``preset_pt`` dataset for replaying pre-tokenized ``.pt`` batches.
	16	+
	17	+Each entry is either a stacked batch dict ``{key: (B, S)-Tensor}`` or a list
	18	+of per-rank dicts. The loader flattens both forms into per-sample rows so the
	19	+standard ``DataLoader`` can batch them again with the trainer collator.
	20	+"""
	21	+import logging
	22	+from typing import Any, Dict, List
	23	+
	24	+import torch
	25	+from torch.utils.data import Dataset
	26	+
	27	+from hyper_parallel.data.registry import DATASET_REGISTRY
	28	+
	29	+
	30	+logger = logging.getLogger(__name__)
	31	+
	32	+
	33	+class PresetPtDataset(Dataset):
	34	+ """Wrap a pre-expanded list of per-sample dicts."""
	35	+
	36	+ def __init__(self, samples: List[Dict[str, Any]]) -> None:
	37	+ self.samples = samples
	38	+
	39	+ def __len__(self) -> int:
	40	+ return len(self.samples)
	41	+
	42	+ def __getitem__(self, idx: int) -> Dict[str, Any]:
	43	+ return self.samples[idx]
	44	+
	45	+
	46	+_OPTIONAL_2D_FIELDS = ("attention_mask", "mm_token_type_ids")
	47	+
	48	+_PIXEL_PAIRS = (("pixel_values", "image_grid_thw"),
	49	+ ("pixel_values_videos", "video_grid_thw"))
	50	+
	51	+
	52	+def _split_pixel_block(
	53	+ b: Dict[str, Any], i: int, batch_size: int, pix_key: str, grid_key: str,
	54	+) -> Dict[str, Any]:
	55	+ """Slice the ``(pix_key, grid_key)`` rows owned by sample ``i``."""
	56	+ pv = b[pix_key]
	57	+ thw = b[grid_key]
	58	+ grids_per_sample = thw.shape[0] // batch_size if thw.dim() == 2 else 0
	59	+ if grids_per_sample == 0:
	60	+ return {}
	61	+ thw_i = thw[i * grids_per_sample:(i + 1) * grids_per_sample].clone()
	62	+ pv_count = int(thw_i.prod(dim=-1).sum().item())
	63	+ offset = sum(
	64	+ int(thw[j].prod(dim=-1).sum().item())
	65	+ for j in range(i * grids_per_sample)
	66	+ )
	67	+ return {
	68	+ pix_key: pv[offset:offset + pv_count].clone(),
	69	+ grid_key: thw_i,
	70	+ }
	71	+
	72	+
	73	+def _slice_position_ids(position_ids: Any, sample_idx: int, batch_size: int) -> Any:
	74	+ """Return one sample's position ids from an LF/HF stacked batch.
	75	+
	76	+ Text Qwen3.5 batches may carry either ``[B, S]`` plain position ids or
	77	+ mRoPE ``[R, B, S]`` ids (``R`` is 3 or 4 in Transformers). Preserve the
	78	+ rotary-rank axis and slice only the batch axis so the trainer can rebuild
	79	+ the original stacked shape in its collate function.
	80	+ """
	81	+ if position_ids.dim() == 2 and position_ids.shape[0] == batch_size:
	82	+ return position_ids[sample_idx].clone()
	83	+ if position_ids.dim() == 3 and position_ids.shape[1] == batch_size:
	84	+ return position_ids[:, sample_idx].clone()
	85	+ if position_ids.dim() == 3 and position_ids.shape[0] == batch_size:
	86	+ return position_ids[sample_idx].clone()
	87	+ raise ValueError(
	88	+ "preset_pt position_ids must be [B, S], [R, B, S], or [B, R, S]; "
	89	+ f"got shape={tuple(position_ids.shape)} with batch_size={batch_size}."
	90	+ )
	91	+
	92	+
	93	+def _expand_batch(b: Dict[str, Any], *, vl: bool) -> List[Dict[str, Any]]:
	94	+ """Split a stacked LM/VL batch dict into per-sample dicts.
	95	+
	96	+ LM samples carry only ``input_ids`` / ``labels`` (plus optionally
	97	+ ``attention_mask``). VL samples additionally carry ``mm_token_type_ids``
	98	+ and ``(pixel_values, image_grid_thw)`` / ``(pixel_values_videos,
	99	+ video_grid_thw)`` pairs sliced according to the per-sample grid product.
	100	+ """
	101	+ ids = b["input_ids"]
	102	+ labels = b["labels"]
	103	+ batch_size = ids.shape[0]
	104	+ out: List[Dict[str, Any]] = []
	105	+ for i in range(batch_size):
	106	+ rec: Dict[str, Any] = {
	107	+ "input_ids": ids[i].clone(),
	108	+ "labels": labels[i].clone(),
	109	+ }
	110	+ if "num_items_in_batch" in b:
	111	+ rec["num_items_in_batch"] = int((rec["labels"] != -100).sum().item())
	112	+ for k in _OPTIONAL_2D_FIELDS:
	113	+ v = b.get(k)
	114	+ if v is not None and v.dim() == 2:
	115	+ rec[k] = v[i].clone()
	116	+ position_ids = b.get("position_ids")
	117	+ if position_ids is not None:
	118	+ rec["position_ids"] = _slice_position_ids(position_ids, i, batch_size)
	119	+ if vl:
	120	+ for pix_key, grid_key in _PIXEL_PAIRS:
	121	+ if b.get(pix_key) is not None and b.get(grid_key) is not None:
	122	+ rec.update(_split_pixel_block(b, i, batch_size, pix_key, grid_key))
	123	+ out.append(rec)
	124	+ return out
	125	+
	126	+
	127	+def _is_vl(batch_entry: Any) -> bool:
	128	+ """Heuristic: VL batches always carry pixel data (image or video)."""
	129	+ if isinstance(batch_entry, list):
	130	+ batch_entry = batch_entry[0] if batch_entry else None
	131	+ return isinstance(batch_entry, dict) and any(
	132	+ pix_key in batch_entry for pix_key, _ in _PIXEL_PAIRS)
	133	+
	134	+
	135	+@DATASET_REGISTRY.register("preset_pt")
	136	+def build_preset_pt(, base: Any, args: Any, *_: Any) -> PresetPtDataset:
	137	+ """Build the preset replay dataset.
	138	+
	139	+ Auto-detects whether the .pt holds VL batches (pixel_values present)
	140	+ or plain LM batches by inspecting the first entry, and dispatches the
	141	+ matching per-sample expansion.
	142	+ """
	143	+ train_path = args.data.train_path
	144	+ if not train_path:
	145	+ raise ValueError("data.train_path is required when data.type='preset_pt'")
	146	+ batches = torch.load(train_path, map_location="cpu", weights_only=False)
	147	+ if not isinstance(batches, list) or not batches:
	148	+ raise ValueError(f"preset_pt expects List, got {type(batches)}")
	149	+
	150	+ vl = _is_vl(batches[0])
	151	+ per_sample: List[Dict[str, Any]] = []
	152	+ for b in batches:
	153	+ if isinstance(b, list):
	154	+ for br in b:
	155	+ per_sample.extend(_expand_batch(br, vl=vl))
	156	+ else:
	157	+ per_sample.extend(_expand_batch(b, vl=vl))
	158	+
	159	+ ds = PresetPtDataset(per_sample)
	160	+ if args.train.max_steps:
	161	+ base.state.max_steps = int(args.train.max_steps)
	162	+ logger.info(
	163	+ "preset_pt dataset (%s): %d samples loaded from %s",
	164	+ "vl" if vl else "lm", len(per_sample), train_path,
	165	+ )
	166	+ return ds

Ahyper_parallel/data/registry.py+108-0

@@ -0,0 +1,108 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Dataset registry — single ``data.type`` name → builder dispatch.
	16	+
	17	+The registry is the only piece of dataset-related code aware of the full
	18	+catalogue of supported formats. Each format lives in its own module and
	19	+registers a builder via ``@DATASET_REGISTRY.register("<name>")``; the trainer
	20	+then calls :func:`build_dataset` with the ``data.type`` from the YAML config
	21	+and never needs to know about specific formats. New formats plug in by adding
	22	+one module-level registration instead of modifying trainer code.
	23	+"""
	24	+from typing import Any, Callable, Dict, Iterator, List
	25	+
	26	+
	27	+BuilderFn = Callable[..., Any]
	28	+
	29	+
	30	+class DatasetRegistry:
	31	+ """Map ``data.type`` names to builder callables.
	32	+
	33	+ Each builder takes keyword arguments ``base`` (the :class:`BaseTrainer`
	34	+ instance), ``args`` (the parsed config) plus any extras the builder
	35	+ needs (``tokenizer``, ``data_transform`` ...). The registry never holds
	36	+ state about a specific run — builders are looked up fresh each time.
	37	+
	38	+ Example:
	39	+ >>> @DATASET_REGISTRY.register("dummy")
	40	+ ... def build_dummy(, base, args, *_):
	41	+ ... return _DummyDataset(...)
	42	+ """
	43	+
	44	+ def __init__(self) -> None:
	45	+ self._builders: Dict[str, BuilderFn] = {}
	46	+
	47	+ def register(self, name: str) -> Callable[[BuilderFn], BuilderFn]:
	48	+ """Decorator that registers ``fn`` under ``name``.
	49	+
	50	+ Raises:
	51	+ ValueError: When ``name`` is already registered. Re-registration
	52	+ is treated as a programming error (typo / duplicate import)
	53	+ rather than silently overwritten.
	54	+ """
	55	+ if not name:
	56	+ raise ValueError("Dataset registry name must be a non-empty string")
	57	+
	58	+ def decorator(fn: BuilderFn) -> BuilderFn:
	59	+ if name in self._builders:
	60	+ raise ValueError(
	61	+ f"Dataset type '{name}' is already registered; "
	62	+ f"check for duplicate registrations / circular imports"
	63	+ )
	64	+ self._builders[name] = fn
	65	+ return fn
	66	+
	67	+ return decorator
	68	+
	69	+ def get(self, name: str) -> BuilderFn:
	70	+ """Return the builder for ``name`` or raise with a helpful message."""
	71	+ try:
	72	+ return self._builders[name]
	73	+ except KeyError as exc:
	74	+ available = sorted(self._builders)
	75	+ raise ValueError(
	76	+ f"Unknown data.type '{name}'. Available: {available}. "
	77	+ f"Add a new format by registering a builder under "
	78	+ f"hyper_parallel.data."
	79	+ ) from exc
	80	+
	81	+ def names(self) -> List[str]:
	82	+ """Return registered names in deterministic (sorted) order."""
	83	+ return sorted(self._builders)
	84	+
	85	+ def __contains__(self, name: str) -> bool:
	86	+ return name in self._builders
	87	+
	88	+ def __iter__(self) -> Iterator[str]:
	89	+ return iter(self.names())
	90	+
	91	+
	92	+DATASET_REGISTRY = DatasetRegistry()
	93	+
	94	+
	95	+def build_dataset(name: str, **kwargs: Any) -> Any:
	96	+ """Build a dataset by ``data.type`` name.
	97	+
	98	+ Args:
	99	+ name: Dataset type identifier — must match a registered builder.
	100	+ **kwargs: Forwarded to the builder. The trainer passes ``base``
	101	+ (the :class:`BaseTrainer` instance) and ``args`` (the parsed
	102	+ config) plus, when available, ``tokenizer`` / ``data_transform``.
	103	+
	104	+ Returns:
	105	+ A ``torch.utils.data.Dataset``-compatible object suitable for the
	106	+ trainer's distributed sampler + ``StatefulDataLoader``.
	107	+ """
	108	+ return DATASET_REGISTRY.get(name)(**kwargs)

Ahyper_parallel/data/vl_dummy.py+181-0

@@ -0,0 +1,181 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""``vl_dummy`` — synthetic VL dataset for native Qwen3-VL smoke tests.
	16	+
	17	+Each sample's pixel grid, image-token slice and trailing token stream are
	18	+seeded by ``train.seed + idx`` so 1-card vs N-card runs see identical content
	19	+after the distributed sampler interleaves them.
	20	+"""
	21	+import logging
	22	+from typing import Any
	23	+
	24	+import torch
	25	+from torch.utils.data import Dataset
	26	+
	27	+from hyper_parallel.data.registry import DATASET_REGISTRY
	28	+
	29	+
	30	+logger = logging.getLogger(__name__)
	31	+
	32	+
	33	+class DummyVLDataset(Dataset):
	34	+ """Synthetic multimodal samples with stable shapes per index."""
	35	+
	36	+ def __init__(
	37	+ self,
	38	+ num_samples: int,
	39	+ seq_length: int,
	40	+ grid_t: int,
	41	+ grid_h: int,
	42	+ grid_w: int,
	43	+ row_width: int,
	44	+ image_tokens: int,
	45	+ image_token_id: int,
	46	+ base_seed: int,
	47	+ video_token_id: int = 151656,
	48	+ is_video: bool = False,
	49	+ ) -> None:
	50	+ self.num_samples = num_samples
	51	+ self.seq_length = seq_length
	52	+ self.grid_t = grid_t
	53	+ self.grid_h = grid_h
	54	+ self.grid_w = grid_w
	55	+ self.row_width = row_width
	56	+ self.image_tokens = image_tokens
	57	+ self.image_token_id = image_token_id
	58	+ self.base_seed = base_seed
	59	+ self.video_token_id = video_token_id
	60	+ self.is_video = is_video
	61	+
	62	+ def __len__(self) -> int:
	63	+ return self.num_samples
	64	+
	65	+ def __getitem__(self, idx: int):
	66	+ g = torch.Generator().manual_seed(self.base_seed + idx)
	67	+ if self.is_video:
	68	+ return self._video_item(idx, g)
	69	+ pixel_values = torch.randn(
	70	+ self.grid_t * self.grid_h * self.grid_w, self.row_width,
	71	+ generator=g, dtype=torch.float32,
	72	+ )
	73	+ input_ids = torch.full((self.seq_length,), 100, dtype=torch.long)
	74	+ input_ids[0] = 151643
	75	+ input_ids[1: 1 + self.image_tokens] = self.image_token_id
	76	+ tail = torch.arange(
	77	+ 200 + idx % 17,
	78	+ 200 + idx % 17 + self.seq_length - 1 - self.image_tokens,
	79	+ dtype=torch.long,
	80	+ )
	81	+ input_ids[1 + self.image_tokens:] = tail
	82	+ labels = input_ids.clone()
	83	+ mm_token_type_ids = torch.zeros(self.seq_length, dtype=torch.int32)
	84	+ mm_token_type_ids[1: 1 + self.image_tokens] = 1
	85	+ return {
	86	+ "input_ids": input_ids,
	87	+ "labels": labels,
	88	+ "attention_mask": torch.ones(self.seq_length, dtype=torch.long),
	89	+ "mm_token_type_ids": mm_token_type_ids,
	90	+ "pixel_values": pixel_values,
	91	+ "image_grid_thw": torch.tensor(
	92	+ [self.grid_t, self.grid_h, self.grid_w], dtype=torch.long,
	93	+ ),
	94	+ }
	95	+
	96	+ def _video_item(self, idx: int, g: torch.Generator):
	97	+ """Build one deterministic video sample."""
	98	+ # A video lays each of grid_t frames as its own mm==2 run (a text
	99	+ # separator between frames) so get_rope_index consumes one per-frame
	100	+ # [1,h,w] grid per run; an image is a single contiguous run.
	101	+ tokens_per_frame = self.image_tokens // self.grid_t
	102	+ pixel_values_videos = torch.randn(
	103	+ self.grid_t * self.grid_h * self.grid_w, self.row_width,
	104	+ generator=g, dtype=torch.float32,
	105	+ )
	106	+ input_ids = torch.full((self.seq_length,), 100, dtype=torch.long)
	107	+ input_ids[0] = 151643
	108	+ mm_token_type_ids = torch.zeros(self.seq_length, dtype=torch.int32)
	109	+ pos = 1
	110	+ frame_idx = 0
	111	+ while frame_idx < self.grid_t:
	112	+ input_ids[pos] = 150 # text separator → split per-frame runs
	113	+ pos += 1
	114	+ input_ids[pos: pos + tokens_per_frame] = self.video_token_id
	115	+ mm_token_type_ids[pos: pos + tokens_per_frame] = 2
	116	+ pos += tokens_per_frame
	117	+ frame_idx += 1
	118	+ tail = torch.arange(
	119	+ 200 + idx % 17, 200 + idx % 17 + self.seq_length - pos,
	120	+ dtype=torch.long,
	121	+ )
	122	+ input_ids[pos:] = tail
	123	+ labels = input_ids.clone()
	124	+ return {
	125	+ "input_ids": input_ids,
	126	+ "labels": labels,
	127	+ "attention_mask": torch.ones(self.seq_length, dtype=torch.long),
	128	+ "mm_token_type_ids": mm_token_type_ids,
	129	+ "pixel_values_videos": pixel_values_videos,
	130	+ "video_grid_thw": torch.tensor(
	131	+ [self.grid_t, self.grid_h, self.grid_w], dtype=torch.long,
	132	+ ),
	133	+ }
	134	+
	135	+
	136	+@DATASET_REGISTRY.register("vl_dummy")
	137	+def build_vl_dummy(, base: Any, args: Any, *_: Any) -> DummyVLDataset:
	138	+ """Build the deterministic VL dummy dataset.
	139	+
	140	+ Pulls ``patch_size`` / ``temporal_patch_size`` / ``in_channels`` /
	141	+ ``spatial_merge_size`` from ``model.config_overrides.vision_config`` so
	142	+ the row width matches the Qwen3-VL vision-tower expectation; falls back
	143	+ to the published defaults when those keys are absent.
	144	+ """
	145	+ max_steps = args.train.max_steps
	146	+ global_bs = args.train.global_batch_size
	147	+ total_samples = max_steps * global_bs
	148	+
	149	+ data_cfg = args.data
	150	+ model_cfg = args.model
	151	+ extra = model_cfg.config_overrides or {}
	152	+ vision_extra = extra.get("vision_config", {}) if isinstance(extra, dict) else {}
	153	+ patch_size = int(vision_extra.get("patch_size", 16))
	154	+ temporal_patch_size = int(vision_extra.get("temporal_patch_size", 2))
	155	+ in_channels = int(vision_extra.get("in_channels", 3))
	156	+ spatial_merge = int(vision_extra.get("spatial_merge_size", 2))
	157	+ grid_t = int(data_cfg.vl_grid_t)
	158	+ grid_h = int(data_cfg.vl_grid_h)
	159	+ grid_w = int(data_cfg.vl_grid_w)
	160	+ image_token_id = int(data_cfg.image_token_id)
	161	+ video_token_id = int(data_cfg.video_token_id)
	162	+ is_video = bool(data_cfg.vl_video)
	163	+ image_tokens = grid_t * grid_h * grid_w // (spatial_merge ** 2)
	164	+ tokens_per_frame = grid_h * grid_w // (spatial_merge ** 2)
	165	+ row_width = in_channels * temporal_patch_size * patch_size * patch_size
	166	+ min_len = grid_t * (tokens_per_frame + 1) + 2 if is_video else image_tokens + 4
	167	+ seq_len = max(int(data_cfg.max_seq_len), min_len)
	168	+ base_seed = int(args.train.seed)
	169	+
	170	+ ds = DummyVLDataset(
	171	+ num_samples=total_samples, seq_length=seq_len,
	172	+ grid_t=grid_t, grid_h=grid_h, grid_w=grid_w, row_width=row_width,
	173	+ image_tokens=image_tokens, image_token_id=image_token_id, base_seed=base_seed,
	174	+ video_token_id=video_token_id, is_video=is_video,
	175	+ )
	176	+ base.state.max_steps = max_steps
	177	+ logger.info(
	178	+ "VL dummy dataset created: samples=%d seq_len=%d grid=(%d,%d,%d) image_tokens=%d",
	179	+ total_samples, seq_len, grid_t, grid_h, grid_w, image_tokens,
	180	+ )
	181	+ return ds

Mhyper_parallel/models/qwen3_5/__init__.py+69-14

	@@ -22,19 +22,20 @@ via YAML ``model.*`` (universal fields) or ``model.config_overrides`` for
22	Qwen-specific ones; everything else falls through to the 0.8B-Base	22	Qwen-specific ones; everything else falls through to the 0.8B-Base
23	defaults baked into :class:`Qwen3_5Config`.	23	defaults baked into :class:`Qwen3_5Config`.
24	"""	24	"""
25	-	25	+import json
26	-__all__ = [	26	+import os
27	- "Qwen3_5Config",	27	+from dataclasses import fields as _dataclass_fields
28	- "Qwen3_5Decoder",
29	- "Qwen3_5ForCausalLM",
30	-]
31		28
32	from hyper_parallel.models.qwen3_5.model import (	29	from hyper_parallel.models.qwen3_5.model import (
33	Qwen3_5Config,	30	Qwen3_5Config,
34	Qwen3_5Decoder,	31	Qwen3_5Decoder,
35	Qwen3_5ForCausalLM,	32	Qwen3_5ForCausalLM,
36	)	33	)
37	-from hyper_parallel.models.qwen3_5.parallelize import parallelize_qwen3_5	34	+from hyper_parallel.models.qwen3_5.parallelize import (
		35	+ parallelize_qwen3_5,
		36	+ pipeline_qwen3_5_for_trainer,
		37	+ qwen3_5_tp_load_transforms,
		38	+)
38	from hyper_parallel.models.qwen3_5.state_dict import (	39	from hyper_parallel.models.qwen3_5.state_dict import (
39	Qwen3_5StateDictAdapter,	40	Qwen3_5StateDictAdapter,
40	)	41	)
	@@ -54,6 +55,38 @@ _UNIVERSAL_FIELDS = (
54	"max_position_embeddings",	55	"max_position_embeddings",
55	)	56	)
56		57
		58	+def _ckpt_config_kwargs(model_cfg) -> dict:
		59	+ """Architecture kwargs read from the checkpoint's ``config.json``.
		60	+
		61	+ Dense hub checkpoints use the composite layout (architecture nested
		62	+ under ``text_config``, ``tie_word_embeddings`` at the top level, rope
		63	+ knobs under ``rope_parameters``); a flat layout is accepted as
		64	+ fallback. Keys are filtered to :class:`Qwen3_5Config` fields so
		65	+ hub-only keys (``hidden_act``, ``mtp_num_hidden_layers`` ...) never
		66	+ reach the dataclass constructor. Without this layer, any dense sibling
		67	+ whose geometry differs from the baked 0.8B-Base defaults would
		68	+ silently build the wrong architecture unless every field is spelled
		69	+ out in YAML.
		70	+ """
		71	+ weights_path = model_cfg.weights_path
		72	+ if not weights_path:
		73	+ return {}
		74	+ cfg_path = os.path.join(weights_path, "config.json")
		75	+ if not os.path.isfile(cfg_path):
		76	+ return {}
		77	+ with open(cfg_path, "r", encoding="utf-8") as handle:
		78	+ raw = json.load(handle)
		79	+ raw_text = dict(raw.get("text_config", raw))
		80	+ rope = raw_text.pop("rope_parameters", None) or {}
		81	+ field_names = {f.name for f in _dataclass_fields(Qwen3_5Config)}
		82	+ kwargs = {k: v for k, v in raw_text.items() if k in field_names}
		83	+ for key in ("rope_theta", "partial_rotary_factor", "mrope_section"):
		84	+ if key in rope:
		85	+ kwargs[key] = rope[key]
		86	+ if "tie_word_embeddings" in raw:
		87	+ kwargs["tie_word_embeddings"] = raw["tie_word_embeddings"]
		88	+ return kwargs
		89	+
57		90
58	def _resolve_overrides(model_cfg) -> dict:	91	def _resolve_overrides(model_cfg) -> dict:
59	"""Merge universal ``ModelConfig`` fields with ``config_overrides`` dict.	92	"""Merge universal ``ModelConfig`` fields with ``config_overrides`` dict.
	@@ -63,20 +96,23 @@ def _resolve_overrides(model_cfg) -> dict:
63	on a per-config basis.	96	on a per-config basis.
64	"""	97	"""
65	overrides: dict = {}	98	overrides: dict = {}
66	- if model_cfg is None:
67	- return overrides
68	for field in _UNIVERSAL_FIELDS:	99	for field in _UNIVERSAL_FIELDS:
69	val = getattr(model_cfg, field, None)	100	val = getattr(model_cfg, field, None)
70	if val is not None:	101	if val is not None:
71	overrides[field] = val	102	overrides[field] = val
72	- extra = getattr(model_cfg, "config_overrides", None)	103	+ extra = model_cfg.config_overrides
73	if isinstance(extra, dict):	104	if isinstance(extra, dict):
74	overrides.update(extra)	105	overrides.update(extra)
75	return overrides	106	return overrides
76		107
77		108
78	def _build(cfg) -> Qwen3_5ForCausalLM:	109	def _build(cfg) -> Qwen3_5ForCausalLM:
79	- """Construct the dense Qwen3.5 model from a ``HyperTrainerConfig``."""	110	+ """Construct the dense Qwen3.5 model from a ``HyperTrainerConfig``.
		111	+
		112	+ Precedence (low → high): dataclass defaults < checkpoint
		113	+ ``config.json`` < universal ``ModelConfig`` fields <
		114	+ ``model.config_overrides``.
		115	+ """
80	return Qwen3_5ForCausalLM(_build_config(cfg))	116	return Qwen3_5ForCausalLM(_build_config(cfg))
81		117
82		118
	@@ -85,10 +121,21 @@ def _build_config(cfg) -> Qwen3_5Config:
85		121
86	Returns the model-specific config object without building the full	122	Returns the model-specific config object without building the full
87	model. Used by the SAPP-ND Hyper YAML parser for memory estimation.	123	model. Used by the SAPP-ND Hyper YAML parser for memory estimation.
88	- """
89	- overrides = _resolve_overrides(getattr(cfg, "model", None))
90	- return Qwen3_5Config(**overrides) if overrides else Qwen3_5Config()
91		124
		125	+ Precedence (low → high): dataclass defaults < checkpoint
		126	+ ``config.json`` < universal ``ModelConfig`` fields <
		127	+ ``model.config_overrides``.
		128	+ """
		129	+ model_cfg = cfg.model
		130	+ overrides = _ckpt_config_kwargs(model_cfg)
		131	+ overrides.update(_resolve_overrides(model_cfg))
		132	+ # A reduced num_hidden_layers must also clip config.json's layer_types
		133	+ # or Qwen3_5Config.__post_init__ rejects the length mismatch.
		134	+ n_layers = overrides.get("num_hidden_layers")
		135	+ layer_types = overrides.get("layer_types")
		136	+ if n_layers is not None and layer_types is not None and len(layer_types) > n_layers:
		137	+ overrides["layer_types"] = layer_types[:n_layers]
		138	+ return Qwen3_5Config(**overrides) if overrides else Qwen3_5Config()
92		139
93	register_spec(	140	register_spec(
94	"qwen3_5",	141	"qwen3_5",
	@@ -96,6 +143,14 @@ register_spec(
96	name="qwen3_5",	143	name="qwen3_5",
97	build_model_fn=_build,	144	build_model_fn=_build,
98	parallelize_fn=parallelize_qwen3_5,	145	parallelize_fn=parallelize_qwen3_5,
		146	+ pipelining_fn=pipeline_qwen3_5_for_trainer,
99	state_dict_adapter=Qwen3_5StateDictAdapter,	147	state_dict_adapter=Qwen3_5StateDictAdapter,
		148	+ tp_load_transform_fn=qwen3_5_tp_load_transforms,
100	),	149	),
101	)	150	)
		151	+
		152	+__all__ = [
		153	+ "Qwen3_5Config",
		154	+ "Qwen3_5Decoder",
		155	+ "Qwen3_5ForCausalLM",
		156	+]

Mhyper_parallel/models/qwen3_5/checkpoint.py+38-6

	@@ -28,8 +28,15 @@ Key remapping (identical to qwen3_5_moe except NO expert unpacking)::
28	model.visual.* → (silently dropped)	28	model.visual.* → (silently dropped)
29	model.mtp.* → (silently dropped)	29	model.mtp.* → (silently dropped)
30		30
31	-When ``tie_word_embeddings=True`` (default for 0.8B-Base), HF drops	31	+Linear-attention layers keep the upstream combined
32	-``lm_head.weight`` from the safetensors index; hyper handles this	32	+``linear_attn.in_proj_qkv.weight`` of shape
		33	+``(key_dim * 2 + value_dim, hidden_size)`` (block layout ``[Q \| K \| V]``
		34	+along axis 0). Tensor parallelism slices that fused tensor by block at
		35	+parallelize/load time so single-card and sharded runs use the same source
		36	+checkpoint layout.
		37	+
		38	+When ``tie_word_embeddings=True`` (default for 0.8B-Base), some checkpoints
		39	+omit ``lm_head.weight`` from the safetensors index; hyper handles this
33	transparently because :class:`Qwen3_5ForCausalLM.__init__` re-binds	40	transparently because :class:`Qwen3_5ForCausalLM.__init__` re-binds
34	``self.lm_head.weight = self.embed_tokens.weight`` at construction time.	41	``self.lm_head.weight = self.embed_tokens.weight`` at construction time.
35	"""	42	"""
	@@ -48,21 +55,21 @@ logger = logging.getLogger(__name__)
48		55
49		56
50	def _remap_simple_key(hf_key: str, max_layer: int) -> Optional[str]:	57	def _remap_simple_key(hf_key: str, max_layer: int) -> Optional[str]:
51	- """Return hyper key for an HF dense-Qwen3.5 key, or None to skip.	58	+ """Return hyper key for a dense-Qwen3.5 checkpoint key, or None to skip.
52		59
53	``max_layer`` is ``num_hidden_layers - 1``: any layer index beyond it	60	``max_layer`` is ``num_hidden_layers - 1``: any layer index beyond it
54	is silently dropped (so hyper can use a smaller truncated model than	61	is silently dropped (so hyper can use a smaller truncated model than
55	the full checkpoint without erroring).	62	the full checkpoint without erroring).
56	"""	63	"""
57	# Vision tower & MTP head — text-only model, drop silently.	64	# Vision tower & MTP head — text-only model, drop silently.
58	- # HF 0.8B-Base stores MTP params at top-level ``mtp.*`` (not	65	+ # 0.8B-Base stores MTP params at top-level ``mtp.*`` (not
59	# ``model.mtp.*``), so we match both forms.	66	# ``model.mtp.*``), so we match both forms.
60	if hf_key.startswith("model.visual."):	67	if hf_key.startswith("model.visual."):
61	return None	68	return None
62	if hf_key.startswith("model.mtp.") or hf_key.startswith("mtp."):	69	if hf_key.startswith("model.mtp.") or hf_key.startswith("mtp."):
63	return None	70	return None
64		71
65	- # lm_head sits at root in HF, same in hyper.	72	+ # lm_head sits at root in the source checkpoint, same in hyper.
66	if hf_key == "lm_head.weight":	73	if hf_key == "lm_head.weight":
67	return "lm_head.weight"	74	return "lm_head.weight"
68		75
	@@ -82,11 +89,28 @@ def _remap_simple_key(hf_key: str, max_layer: int) -> Optional[str]:
82	logger.debug("Unmapped HF key dropped: %s", hf_key)	89	logger.debug("Unmapped HF key dropped: %s", hf_key)
83	return None	90	return None
84		91
		92	+def _maybe_split_in_proj_qkv(
		93	+ hyper_sd: Dict[str, torch.Tensor],
		94	+ hyper_key: str,
		95	+ tensor: torch.Tensor,
		96	+ key_dim: int,
		97	+) -> None:
		98	+ """Insert ``tensor`` into ``hyper_sd``.
		99	+
		100	+ The legacy name is kept for callers/tests that exercise the loader helper.
		101	+ Dense Qwen3.5 now keeps the checkpoint's fused ``in_proj_qkv`` parameter;
		102	+ TP blockwise slicing is handled later by :func:`qwen3_5_tp_load_transforms`.
		103	+ """
		104	+ del key_dim
		105	+ hyper_sd[hyper_key] = tensor
		106	+
85		107
86	def load_hf_qwen3_5_state_dict(	108	def load_hf_qwen3_5_state_dict(
87	weights_path: str,	109	weights_path: str,
88	num_hidden_layers: int,	110	num_hidden_layers: int,
89	dtype: Optional[torch.dtype] = None,	111	dtype: Optional[torch.dtype] = None,
		112	+ *,
		113	+ linear_key_dim: Optional[int] = None,
90	) -> Dict[str, torch.Tensor]:	114	) -> Dict[str, torch.Tensor]:
91	"""Load a dense Qwen3.5 safetensors checkpoint into hyper state_dict.	115	"""Load a dense Qwen3.5 safetensors checkpoint into hyper state_dict.
92		116
	@@ -97,6 +121,8 @@ def load_hf_qwen3_5_state_dict(
97	``num_hidden_layers`` text-decoder layers (later layers are	121	``num_hidden_layers`` text-decoder layers (later layers are
98	silently dropped).	122	silently dropped).
99	dtype: Optional cast applied to every returned tensor.	123	dtype: Optional cast applied to every returned tensor.
		124	+ linear_key_dim: Kept for API compatibility with older split-QKV
		125	+ loaders. Dense Qwen3.5 keeps the combined checkpoint key unchanged.
100		126
101	Returns:	127	Returns:
102	Dict keyed by hyper module paths, values on CPU.	128	Dict keyed by hyper module paths, values on CPU.
	@@ -143,7 +169,13 @@ def load_hf_qwen3_5_state_dict(
143	if mapped is None:	169	if mapped is None:
144	skipped += 1	170	skipped += 1
145	continue	171	continue
146	- hyper_sd[mapped] = _cast(f.get_tensor(hf_key))	172	+ tensor = _cast(f.get_tensor(hf_key))
		173	+ if linear_key_dim is not None:
		174	+ _maybe_split_in_proj_qkv(
		175	+ hyper_sd, mapped, tensor, linear_key_dim,
		176	+ )
		177	+ else:
		178	+ hyper_sd[mapped] = tensor
147		179
148	# ``tie_word_embeddings=True`` checkpoints drop ``lm_head.weight``;	180	# ``tie_word_embeddings=True`` checkpoints drop ``lm_head.weight``;
149	# ``fully_shard`` breaks the Python-level tie, so synthesize the tensor	181	# ``fully_shard`` breaks the Python-level tie, so synthesize the tensor

Mhyper_parallel/models/qwen3_5/model.py+433-43

	@@ -14,12 +14,12 @@
14	# ============================================================================	14	# ============================================================================
15	"""Qwen3.5 (dense) — hybrid linear-attention + gated-attention, no MoE.	15	"""Qwen3.5 (dense) — hybrid linear-attention + gated-attention, no MoE.
16		16
17	-Mirrors the upstream ``Qwen3_5ForConditionalGeneration`` text architecture	17	+Implements the ``Qwen3_5ForConditionalGeneration`` text architecture
18	(``model_type='qwen3_5'``, e.g. ``Qwen/Qwen3.5-0.8B-Base``):	18	(``model_type='qwen3_5'``, e.g. ``Qwen/Qwen3.5-0.8B-Base``):
19		19
20	- N-layer text decoder, ``layer_types=[linear, linear, linear, full] * (N/4)``	20	- N-layer text decoder, ``layer_types=[linear, linear, linear, full] * (N/4)``
21	- ``linear_attention`` layers use :class:`GatedDeltaNet`	21	- ``linear_attention`` layers use :class:`GatedDeltaNet`
22	- - ``full_attention`` layers use :class:`GroupQueryAttention` with	22	+ - ``full_attention`` layers use :class:`Qwen3_5Attention` with
23	``attn_output_gate=True`` and ``partial_rotary_factor=0.25``	23	``attn_output_gate=True`` and ``partial_rotary_factor=0.25``
24	- Every layer's MLP is a plain :class:`SwiGLUMLP` (gate / up / down).	24	- Every layer's MLP is a plain :class:`SwiGLUMLP` (gate / up / down).
25	No MoE, no shared-expert, no routing — this is what distinguishes	25	No MoE, no shared-expert, no routing — this is what distinguishes

Mhyper_parallel/models/qwen3_5/parallelize.py+1300-53

Mhyper_parallel/models/qwen3_5/state_dict.py+10-2

	@@ -20,7 +20,7 @@ can stay generic. Registered on the ``qwen3_5`` ModelSpec.
20		20
21	__all__ = ["Qwen3_5StateDictAdapter"]	21	__all__ = ["Qwen3_5StateDictAdapter"]
22		22
23	-# pylint: disable=C0103 # HF transformers class-name convention (Qwen3_5*)	23	+# pylint: disable=C0103 # Qwen class-name convention (Qwen3_5*)
24	from typing import Dict, Optional	24	from typing import Dict, Optional
25		25
26	import torch	26	import torch
	@@ -37,11 +37,19 @@ class Qwen3_5StateDictAdapter:
37	model_config,	37	model_config,
38	dtype: Optional[torch.dtype] = None,	38	dtype: Optional[torch.dtype] = None,
39	) -> Dict[str, torch.Tensor]:	39	) -> Dict[str, torch.Tensor]:
40	- """Read an HF safetensors checkpoint and return a hyper-named state dict."""	40	+ """Read an HF safetensors checkpoint and return a hyper-named state dict.
		41	+
		42	+ The dense model keeps the upstream combined
		43	+ ``linear_attn.in_proj_qkv.weight`` key. ``linear_key_dim`` is still
		44	+ passed for compatibility with older loader call sites but no split is
		45	+ performed here; TP-specific block slicing happens in ``parallelize.py``.
		46	+ """
		47	+ linear_key_dim = model_config.linear_key_head_dim * model_config.linear_num_key_heads
41	return load_hf_qwen3_5_state_dict(	48	return load_hf_qwen3_5_state_dict(
42	weights_path,	49	weights_path,
43	num_hidden_layers=model_config.num_hidden_layers,	50	num_hidden_layers=model_config.num_hidden_layers,
44	dtype=dtype,	51	dtype=dtype,
		52	+ linear_key_dim=linear_key_dim,
45	)	53	)
46		54
47	@staticmethod	55	@staticmethod

Mhyper_parallel/models/qwen3_5_moe/__init__.py+171-15

	@@ -20,14 +20,27 @@ geometry, mRoPE section split, linear-attention dims and ``layer_types``
20	— go through ``cfg.model.config_overrides``. Anything not provided falls	20	— go through ``cfg.model.config_overrides``. Anything not provided falls
21	through to the Qwen3.5-35B-A3B defaults below.	21	through to the Qwen3.5-35B-A3B defaults below.
22	"""	22	"""
		23	+import json
		24	+import os
		25	+from dataclasses import fields as _dataclass_fields
		26	+
23	from hyper_parallel.models.qwen3_5_moe.model import (	27	from hyper_parallel.models.qwen3_5_moe.model import (
24	Qwen3_5MoeConfig,	28	Qwen3_5MoeConfig,
25	Qwen3_5MoeForCausalLM,	29	Qwen3_5MoeForCausalLM,
26	)	30	)
27	-from hyper_parallel.models.qwen3_5_moe.parallelize import parallelize_qwen3_5_moe	31	+from hyper_parallel.models.qwen3_5_moe.model_vl import (
		32	+ Qwen3_5MoeVLConfig,
		33	+ Qwen3_5MoeVLForConditionalGeneration,
		34	+)
		35	+from hyper_parallel.models.qwen3_5_moe.parallelize import (
		36	+ parallelize_qwen3_5_moe,
		37	+ pipeline_qwen3_5_moe_for_trainer,
		38	+ qwen3_5_moe_tp_load_transforms,
		39	+)
28	from hyper_parallel.models.qwen3_5_moe.state_dict import (	40	from hyper_parallel.models.qwen3_5_moe.state_dict import (
29	Qwen3_5MoeStateDictAdapter,	41	Qwen3_5MoeStateDictAdapter,
30	)	42	)
		43	+from hyper_parallel.models.qwen3_vl_vision import Qwen3VLMoeVisionConfig
31	from hyper_parallel.models.spec import ModelSpec, register_spec	44	from hyper_parallel.models.spec import ModelSpec, register_spec
32		45
33	# Default Qwen3.5-35B-A3B architecture values taken from	46	# Default Qwen3.5-35B-A3B architecture values taken from
	@@ -60,15 +73,16 @@ _DEFAULTS = {
60	"num_experts_per_tok": 8,	73	"num_experts_per_tok": 8,
61	"moe_intermediate_size": 512,	74	"moe_intermediate_size": 512,
62	"shared_expert_intermediate_size": 512,	75	"shared_expert_intermediate_size": 512,
		76	+ "mtp_loss_weight": 0.0,
63	}	77	}
64		78
65	# Universal transformer fields exposed on ``ModelConfig`` that map straight	79	# Universal transformer fields exposed on ``ModelConfig`` that map straight
66	# onto :class:`Qwen3_5MoeConfig`. All other architecture fields go through	80	# onto :class:`Qwen3_5MoeConfig`. All other architecture fields go through
67	-# ``model.config_overrides``.	81	+# ``model.config_overrides``. (``intermediate_size`` is deliberately absent:
		82	+# the MoE config has no such field — experts use ``moe_intermediate_size``.)
68	_UNIVERSAL_FIELDS = (	83	_UNIVERSAL_FIELDS = (
69	"vocab_size",	84	"vocab_size",
70	"hidden_size",	85	"hidden_size",
71	- "intermediate_size",
72	"num_hidden_layers",	86	"num_hidden_layers",
73	"num_attention_heads",	87	"num_attention_heads",
74	"num_key_value_heads",	88	"num_key_value_heads",
	@@ -76,34 +90,175 @@ _UNIVERSAL_FIELDS = (
76	)	90	)
77		91
78		92
		93	+def _ckpt_text_kwargs(raw: dict) -> dict:
		94	+ """Text-config kwargs extracted from a checkpoint's parsed ``config.json``.
		95	+
		96	+ Handles the composite hub layout (architecture under ``text_config``,
		97	+ ``tie_word_embeddings`` at the top level, rope knobs nested under
		98	+ ``rope_parameters``) with a flat fallback. Keys are filtered to
		99	+ :class:`Qwen3_5MoeConfig` fields so hub-only keys (``hidden_act``,
		100	+ ``mtp_num_hidden_layers`` ...) cannot reach the dataclass constructor.
		101	+ """
		102	+ if not raw:
		103	+ return {}
		104	+ raw_text = dict(raw.get("text_config", raw))
		105	+ rope = raw_text.pop("rope_parameters", None) or {}
		106	+ field_names = {f.name for f in _dataclass_fields(Qwen3_5MoeConfig)}
		107	+ kwargs = {k: v for k, v in raw_text.items() if k in field_names}
		108	+ for key in ("rope_theta", "partial_rotary_factor", "mrope_section"):
		109	+ if key in rope:
		110	+ kwargs[key] = rope[key]
		111	+ if "tie_word_embeddings" in raw:
		112	+ kwargs["tie_word_embeddings"] = raw["tie_word_embeddings"]
		113	+ return kwargs
		114	+
		115	+
		116	+def _truncate_layer_types(kwargs: dict) -> None:
		117	+ """Normalize ``layer_types`` to ``num_hidden_layers`` after all merges.
		118	+
		119	+ A reduced num_hidden_layers (tests shrink the real 40-layer config) must
		120	+ also truncate config.json's layer_types to the first N layers — matching
		121	+ the checkpoint loader's max_layer gate — else Qwen3_5MoeConfig.__post_init__
		122	+ rejects the length mismatch. Some reduced checkpoints store a 1-layer
		123	+ config while still carrying more layer weights; when a caller asks for more
		124	+ layers than the config lists, regenerate the standard Qwen3.5 layer pattern.
		125	+ """
		126	+ n_layers = kwargs.get("num_hidden_layers")
		127	+ layer_types = kwargs.get("layer_types")
		128	+ if n_layers is None or layer_types is None:
		129	+ return
		130	+ if len(layer_types) > n_layers:
		131	+ kwargs["layer_types"] = layer_types[:n_layers]
		132	+ elif len(layer_types) < n_layers:
		133	+ kwargs["layer_types"] = None
		134	+
		135	+
79	def _resolve_kwargs(model_cfg) -> dict:	136	def _resolve_kwargs(model_cfg) -> dict:
80	"""Build the full kwargs dict for :class:`Qwen3_5MoeConfig`.	137	"""Build the full kwargs dict for :class:`Qwen3_5MoeConfig`.
81		138
82	- Precedence (low → high): ``_DEFAULTS`` < universal ``ModelConfig``	139	+ Precedence (low → high): ``_DEFAULTS`` < checkpoint ``config.json`` <
83	- fields < ``model.config_overrides`` dict. ``config_overrides`` wins so	140	+ universal ``ModelConfig`` fields < ``model.config_overrides`` dict.
84	- callers can clobber any architecture knob without us having to surface	141	+ ``config.json`` keeps a sibling-size checkpoint from silently building
85	- it as a typed field on ``ModelConfig``.	142	+ the 35B-A3B default architecture; ``config_overrides`` wins so callers
		143	+ can clobber any architecture knob without us having to surface it as a
		144	+ typed field on ``ModelConfig``.
86	"""	145	"""
87	kwargs = dict(_DEFAULTS)	146	kwargs = dict(_DEFAULTS)
88	- if model_cfg is None:	147	+ kwargs.update(
89	- return kwargs	148	+ _ckpt_text_kwargs(
		149	+ _read_vl_config_json(model_cfg.weights_path)
		150	+ )
		151	+ )
90	for field in _UNIVERSAL_FIELDS:	152	for field in _UNIVERSAL_FIELDS:
91	val = getattr(model_cfg, field, None)	153	val = getattr(model_cfg, field, None)
92	if val is not None:	154	if val is not None:
93	kwargs[field] = val	155	kwargs[field] = val
94	- extra = getattr(model_cfg, "config_overrides", None)	156	+ extra = model_cfg.config_overrides
95	if isinstance(extra, dict):	157	if isinstance(extra, dict):
96	kwargs.update(extra)	158	kwargs.update(extra)
		159	+ _truncate_layer_types(kwargs)
97	return kwargs	160	return kwargs
98		161
		162	+# Token-id fields that live at the top level of the multimodal config.json.
		163	+_VL_TOKEN_FIELDS = (
		164	+ "image_token_id",
		165	+ "video_token_id",
		166	+ "vision_start_token_id",
		167	+ "vision_end_token_id",
		168	+)
99		169
100	-def _build(cfg) -> Qwen3_5MoeForCausalLM:	170	+
101	- """Translate trainer cfg → Qwen3_5MoeConfig → model.	171	+def _read_vl_config_json(weights_path) -> dict:
		172	+ """Load ``config.json`` from a multimodal checkpoint, else return ``{}``."""
		173	+ if not weights_path:
		174	+ return {}
		175	+ cfg_path = os.path.join(weights_path, "config.json")
		176	+ if not os.path.isfile(cfg_path):
		177	+ return {}
		178	+ with open(cfg_path, "r", encoding="utf-8") as handle:
		179	+ return json.load(handle)
		180	+
		181	+
		182	+def _resolve_vl_kwargs(model_cfg) -> tuple:
		183	+ """Build (text_kwargs, vision_kwargs, token_ids) for the VL composite.
		184	+
		185	+ Precedence (low → high): hp ``_DEFAULTS`` < checkpoint ``config.json`` <
		186	+ universal ``ModelConfig`` fields < ``model.config_overrides``.
		187	+ ``config.json`` is authoritative for the real checkpoint architecture;
		188	+ ``config_overrides`` lets tests shrink it.
		189	+ """
		190	+ raw = _read_vl_config_json(model_cfg.weights_path)
		191	+
		192	+ text_field_names = {f.name for f in _dataclass_fields(Qwen3_5MoeConfig)}
		193	+ vision_field_names = {f.name for f in _dataclass_fields(Qwen3VLMoeVisionConfig)}
		194	+
		195	+ text_kwargs = dict(_DEFAULTS)
		196	+ text_kwargs.update(_ckpt_text_kwargs(raw))
		197	+ for field in _UNIVERSAL_FIELDS:
		198	+ val = getattr(model_cfg, field, None)
		199	+ if val is not None:
		200	+ text_kwargs[field] = val
		201	+
		202	+ # DeepStack disabled by default for the Qwen3.5 ViT (the merged ViT).
		203	+ vision_kwargs = {
		204	+ "deepstack_visual_indexes": [],
		205	+ "_attn_implementation": "eager",
		206	+ }
		207	+ for key, val in raw.get("vision_config", {}).items():
		208	+ if key in vision_field_names:
		209	+ vision_kwargs[key] = val
		210	+
		211	+ token_ids = {key: raw[key] for key in _VL_TOKEN_FIELDS if key in raw}
		212	+
		213	+ extra = model_cfg.config_overrides
		214	+ if isinstance(extra, dict):
		215	+ text_extra = extra.get("text_config", {})
		216	+ vision_extra = extra.get("vision_config", {})
		217	+ text_kwargs.update(
		218	+ {k: v for k, v in extra.items()
		219	+ if k in text_field_names and k not in ("vl",)}
		220	+ )
		221	+ text_kwargs.update({k: v for k, v in text_extra.items() if k in text_field_names})
		222	+ vision_kwargs.update({k: v for k, v in vision_extra.items() if k in vision_field_names})
		223	+ token_ids.update({k: v for k, v in extra.items() if k in _VL_TOKEN_FIELDS})
		224	+
		225	+ _truncate_layer_types(text_kwargs)
		226	+
		227	+ return text_kwargs, vision_kwargs, token_ids
		228	+
		229	+
		230	+def _build_vl(cfg) -> Qwen3_5MoeVLForConditionalGeneration:
		231	+ """Build the multimodal ``Qwen3_5MoeVLForConditionalGeneration``."""
		232	+ text_kwargs, vision_kwargs, token_ids = _resolve_vl_kwargs(cfg.model)
		233	+ vl_config = Qwen3_5MoeVLConfig(
		234	+ text_config=Qwen3_5MoeConfig(**text_kwargs),
		235	+ vision_config=Qwen3VLMoeVisionConfig(**vision_kwargs),
		236	+ **token_ids,
		237	+ )
		238	+ return Qwen3_5MoeVLForConditionalGeneration(vl_config)
		239	+
		240	+
		241	+def _build(cfg):
		242	+ """Translate trainer cfg → model.
102		243
103	A minimal YAML ``model: {name: qwen3_5_moe, weights_path: ...}`` is	244	A minimal YAML ``model: {name: qwen3_5_moe, weights_path: ...}`` is
104	- enough — every other field falls through to the 35B-A3B defaults.	245	+ enough — every other field falls through to the 35B-A3B defaults. Set
		246	+ ``config_overrides.vl: true`` to build the multimodal (vision + text)
		247	+ ``Qwen3_5MoeVLForConditionalGeneration`` instead of the text-only LM.
105	"""	248	"""
106	- return Qwen3_5MoeForCausalLM(_build_config(cfg))	249	+ extra = cfg.model.config_overrides
		250	+ if isinstance(extra, dict) and extra.get("vl", False):
		251	+ return _build_vl(cfg)
		252	+ kwargs = _resolve_kwargs(cfg.model)
		253	+ if kwargs.get("mtp_loss_weight", 0.0) > 0:
		254	+ # Only the VL composite constructs the MTP head; accepting the flag
		255	+ # here would silently train without the MTP loss.
		256	+ raise ValueError(
		257	+ "mtp_loss_weight > 0 requires the VL composite "
		258	+ "(config_overrides.vl: true) — the text-only Qwen3_5MoeForCausalLM "
		259	+ "has no MTP head."
		260	+ )
		261	+ return Qwen3_5MoeForCausalLM(Qwen3_5MoeConfig(**kwargs))
107		262
108		263
109	def _build_config(cfg) -> Qwen3_5MoeConfig:	264	def _build_config(cfg) -> Qwen3_5MoeConfig:
	@@ -114,13 +269,14 @@ def _build_config(cfg) -> Qwen3_5MoeConfig:
114	"""	269	"""
115	return Qwen3_5MoeConfig(**_resolve_kwargs(cfg.model))	270	return Qwen3_5MoeConfig(**_resolve_kwargs(cfg.model))
116		271
117	-
118	register_spec(	272	register_spec(
119	"qwen3_5_moe",	273	"qwen3_5_moe",
120	ModelSpec(	274	ModelSpec(
121	name="qwen3_5_moe",	275	name="qwen3_5_moe",
122	build_model_fn=_build,	276	build_model_fn=_build,
123	parallelize_fn=parallelize_qwen3_5_moe,	277	parallelize_fn=parallelize_qwen3_5_moe,
		278	+ pipelining_fn=pipeline_qwen3_5_moe_for_trainer,
124	state_dict_adapter=Qwen3_5MoeStateDictAdapter,	279	state_dict_adapter=Qwen3_5MoeStateDictAdapter,
		280	+ tp_load_transform_fn=qwen3_5_moe_tp_load_transforms,
125	),	281	),
126	)	282	)