Add WALL-OSS-0.5 native decoder and action processor components#38
Add WALL-OSS-0.5 native decoder and action processor components#38Grape203 wants to merge 11 commits into
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces experimental native PyTorch modules for the WALL-OSS-0.5 model, including configuration overlays, action processors, attention projections, M-RoPE, attention core, decoder layers, and the decoder model skeleton, along with several parity and smoke test scripts to validate the implementation against official references. The review feedback focuses on performance optimizations, suggesting that position embeddings be sliced and cast once at the model level rather than redundantly in each layer, that .expand() be used instead of .repeat() for memory efficiency, that slice indices in M-RoPE be precomputed to avoid splitting tensors on every forward pass, and that redundant mask creation be avoided when the attention mask is null to enable SDPA optimized fast-paths.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| query_states, key_states = self.mrope( | ||
| query_states.contiguous(), | ||
| key_states.contiguous(), | ||
| cos[..., : (cos.size(3) // 2)].contiguous().float(), | ||
| sin[..., : (sin.size(3) // 2)].contiguous().float(), | ||
| ) | ||
|
|
There was a problem hiding this comment.
In WallOSS05JointAttentionNative.forward, slicing, casting, and making cos/sin contiguous on every layer's forward pass is highly redundant (done 36 times per forward pass). We can optimize this by checking if the embeddings are already sliced to head_dim // 2, allowing us to perform this operation once at the model level.
cos, sin = position_embeddings
if cos.size(-1) == self.head_dim:
cos = cos[..., : (self.head_dim // 2)].contiguous().float()
sin = sin[..., : (self.head_dim // 2)].contiguous().float()
query_states, key_states = self.mrope(
query_states.contiguous(),
key_states.contiguous(),
cos,
sin,
)| projection_dtype: torch.dtype | None = None, | ||
| apply_final_norm: bool = True, | ||
| ) -> torch.Tensor: | ||
| for layer in self.layers: |
There was a problem hiding this comment.
In WallOSS05DecoderModelNative.forward, we can slice, cast, and make cos/sin contiguous once before the layer loop, saving 35 redundant operations per forward pass.
cos, sin = position_embeddings
cos = cos[..., : (cos.size(3) // 2)].contiguous().float()
sin = sin[..., : (sin.size(3) // 2)].contiguous().float()
position_embeddings = (cos, sin)
for layer in self.layers:| time_embed = self.time_embed(timestep).to(dtype=noisy_action.dtype) | ||
| action_embed = _linear_forward(self.w1, noisy_action) | ||
|
|
||
| time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) |
There was a problem hiding this comment.
In WallOSS05ActionProcessorNative.step, time_embed is repeated using .repeat(). Using .expand() is much more efficient as it creates a view without copying memory, and it is fully compatible with torch.cat.
| time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) | |
| time_embed = time_embed.unsqueeze(1).expand(-1, action_embed.shape[1], -1) |
| def __init__(self, mrope_section: list[int] | tuple[int, ...]): | ||
| super().__init__() | ||
| section = tuple(int(v) for v in mrope_section) | ||
| if len(section) != 3: | ||
| raise ValueError(f"mrope_section must have 3 entries, got {section}") | ||
| if any(v <= 0 for v in section): | ||
| raise ValueError(f"mrope_section entries must be positive, got {section}") | ||
| self.mrope_section = section |
There was a problem hiding this comment.
In WallOSS05MRoPENative.__init__, we can precompute the slice indices for cos and sin to avoid splitting the entire tensor along the last dimension on every forward pass.
def __init__(self, mrope_section: list[int] | tuple[int, ...]):
super().__init__()
section = tuple(int(v) for v in mrope_section)
if len(section) != 3:
raise ValueError(f"mrope_section must have 3 entries, got {section}")
if any(v <= 0 for v in section):
raise ValueError(f"mrope_section entries must be positive, got {section}")
self.mrope_section = section
sections = list(section) + list(section)
offsets = [0]
for s in sections:
offsets.append(offsets[-1] + s)
self.slices = [(i % 3, offsets[i], offsets[i+1]) for i in range(len(sections))]| cos = cos.float() | ||
| sin = sin.float() | ||
|
|
||
| cos = torch.cat((cos, cos), dim=-1) | ||
| sin = torch.cat((sin, sin), dim=-1) | ||
|
|
||
| mrope_section_doubled = list(self.mrope_section) + list(self.mrope_section) | ||
|
|
||
| cos_split = torch.cat( | ||
| [m[i % 3] for i, m in enumerate(cos.split(mrope_section_doubled, dim=-1))], | ||
| dim=-1, | ||
| ).unsqueeze(2) | ||
| sin_split = torch.cat( | ||
| [m[i % 3] for i, m in enumerate(sin.split(mrope_section_doubled, dim=-1))], | ||
| dim=-1, | ||
| ).unsqueeze(2) |
There was a problem hiding this comment.
In WallOSS05MRoPENative.forward, we can use direct slicing with the precomputed slice indices instead of calling cos.split and list comprehension. This is significantly faster and avoids creating intermediate split tensors.
cos = torch.cat((cos, cos), dim=-1)
sin = torch.cat((sin, sin), dim=-1)
cos_split = torch.cat(
[cos[idx, :, :, start:end] for idx, start, end in self.slices],
dim=-1,
).unsqueeze(2)
sin_split = torch.cat(
[sin[idx, :, :, start:end] for idx, start, end in self.slices],
dim=-1,
).unsqueeze(2)| causal_mask = torch.ones( | ||
| bsz, | ||
| 1, | ||
| 1, | ||
| key_len, | ||
| device=device, | ||
| dtype=dtype, | ||
| ).contiguous() | ||
| causal_mask = causal_mask.to(torch.bool) | ||
|
|
There was a problem hiding this comment.
In WallOSS05AttentionCoreNative._prepare_causal_mask, if attention_mask is None and q_len == 1, creating a mask of all True is redundant and prevents SDPA from using optimized unmasked fast-paths. We should only create the ones mask if causal_mask is not None.
| causal_mask = torch.ones( | |
| bsz, | |
| 1, | |
| 1, | |
| key_len, | |
| device=device, | |
| dtype=dtype, | |
| ).contiguous() | |
| causal_mask = causal_mask.to(torch.bool) | |
| if q_len == 1 and causal_mask is not None: | |
| causal_mask = torch.ones( | |
| bsz, | |
| 1, | |
| 1, | |
| key_len, | |
| device=device, | |
| dtype=dtype, | |
| ).contiguous() | |
| causal_mask = causal_mask.to(torch.bool) |
1 participant
Summary
This PR adds an experimental native WALL-OSS-0.5 decoder/action-processor implementation path.
The branch is rebased on the latest
upstream/mainafter Cosmos3 support was merged. The goal of this PR is to move beyond the previous wrapper-based path and provide native PhyAI modeling components for the WALL-OSS-0.5 decoder and action processor.What is included
WallOSS05NativeConfigActionProcessorexamples/Validation
All native checks passed after rebasing onto latest
upstream/main.Validated scripts include:
run_walloss05_native_action_processor_check.pyrun_walloss05_native_action_processor_parity.pyrun_walloss05_native_sparse_moe_parity.pyrun_walloss05_native_norm_moe_parity.pyrun_walloss05_native_decoder_ffn_parity.pyrun_walloss05_native_attention_projection_parity.pyrun_walloss05_native_mrope_parity.pyrun_walloss05_native_attention_core_parity.pyrun_walloss05_native_joint_attention_parity.pyrun_walloss05_native_decoder_layer_parity.pyrun_walloss05_native_decoder_model_skeleton_check.pyThe current implementation verifies native decoder/action-processor components, layer-0 no-cache /
mot_opt=Truepaths, and 36-layer decoder checkpoint tensor shape alignment.Scope
This PR does not claim full end-to-end WALL-OSS-0.5 runtime integration yet.
Not included in this PR:
Those parts are intended for the next integration stage.