feat: Major optimization and robustness improvements (#31)

This release introduces significant performance optimizations, memory efficiency
improvements, and enhanced system robustness:

🚀 Performance Optimizations:
- Add three new fused CUDA kernels (rope_index, rot_pos_emb, get_window_index)
  for accelerated multimodal preprocessing
- Implement FSDP2 support for distributed training with improved memory efficiency
- Add Torch.compile integration for additional performance gains
- Optimize memory usage: reduce peak allocation from 48GB to 24GB on 8-GPU setup

🔧 System Robustness:
- Fix missing token position inputs in prediction pipeline
- Add type-robust negation operations in RoPE CUDA kernels (half/bfloat16 support)
- Fix dataset root parameter initialization in LeRobot data loader
- Enhanced error handling and input validation across fusion operators

📚 Documentation & Usability:
- Add comprehensive memory usage benchmarks and hardware recommendations
- Update citation format with proper arXiv reference
- Improve training configuration documentation with quick start guide
- Add detailed API documentation for new fusion operators

🛠️ Technical Details:
- Version bump to 1.0.1
- New CUDA kernels: rope_index.cu, rot_pos.cu, window_index.cu
- FSDP2 state dict loading with distribute_tensor support
- Enhanced multimodal RoPE with 3D position encoding
- Window attention optimization for Vision Transformers

Breaking Changes: None - all changes are backward compatible

wall_x/data/load_lerobot_dataset.py

@@ -424,8 +424,9 @@ def load_lerobot_data(
 
     # repo_id = "lerobot/aloha_mobile_cabinet"
     repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
+    root = lerobot_config.get("root", None)
     dataset = LeRobotDataset(
-        repo_id, delta_timestamps=delta_timestamps, video_backend="pyav"
+        repo_id, root=root, delta_timestamps=delta_timestamps, video_backend="pyav"
     )
 
     if rank == 0:

wall_x/fusions/backend.py

@@ -297,3 +297,137 @@ def rope_bwd(
     return backend.rope_bwd(
         grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled
     )
+
+
+def get_rope_index(
+    input_ids: torch.Tensor,
+    image_grid_thw: Optional[torch.Tensor],
+    video_grid_thw: Optional[torch.Tensor],
+    second_per_grid_ts: Optional[torch.Tensor],
+    attention_mask: Optional[torch.Tensor],
+    spatial_merge_size: int,
+    image_token_id: int,
+    video_token_id: int,
+    vision_start_token_id: int,
+    tokens_per_second: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate position indices for multimodal RoPE (Rotary Position Embedding).
+
+    This function computes 3D position indices for text, image, and video tokens
+    to enable proper spatial-temporal position encoding in multimodal transformers.
+
+    Args:
+        input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len]
+        image_grid_thw (torch.Tensor, optional): Image grid specifications of shape [num_images, 3] (T, H, W)
+        video_grid_thw (torch.Tensor, optional): Video grid specifications of shape [num_videos, 3] (T, H, W)
+        second_per_grid_ts (torch.Tensor, optional): Temporal scaling per video grid of shape [num_videos]
+        attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len]
+        spatial_merge_size (int): Spatial dimension merge factor for patch grouping
+        image_token_id (int): Token ID representing image patches
+        video_token_id (int): Token ID representing video frames
+        vision_start_token_id (int): Token ID marking vision sequence start
+        tokens_per_second (float): Temporal scaling factor for video sequences
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - position_ids: 3D position indices of shape [3, batch_size, seq_len]
+            - mrope_deltas: Position deltas for multimodal RoPE of shape [batch_size, 1]
+
+    Note:
+        When both image_grid_thw and video_grid_thw are None, returns standard
+        text-only position indices based on attention_mask or sequence order.
+    """
+    return backend.rope_index(
+        input_ids,
+        image_grid_thw,
+        video_grid_thw,
+        second_per_grid_ts,
+        attention_mask,
+        spatial_merge_size,
+        image_token_id,
+        video_token_id,
+        vision_start_token_id,
+        tokens_per_second,
+    )
+
+
+def rot_pos_emb(
+    inv_freq: torch.Tensor,
+    grid_thw: torch.Tensor,
+    spatial_merge_size: int,
+) -> torch.Tensor:
+    """
+    Compute fused rotary position embeddings for multimodal grids.
+
+    This function efficiently computes rotary position embeddings for spatial-temporal
+    grids using a fused CUDA kernel, supporting both int32 and int64 grid specifications.
+
+    Args:
+        inv_freq (torch.Tensor): Inverse frequencies for RoPE of shape [dim/2]
+                                Must be float32 dtype on CUDA device
+        grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W)
+                                Supports int32 or int64 dtype on CUDA device
+        spatial_merge_size (int): Merge factor for spatial dimensions (must be positive)
+
+    Returns:
+        torch.Tensor: Computed rotary embeddings of shape [total_tokens, dim]
+                     where total_tokens is determined by grid layouts and spatial_merge_size
+
+    Example:
+        >>> inv_freq = torch.randn(64, device='cuda', dtype=torch.float32)  # 128-dim model
+        >>> grids = torch.tensor([[8, 14, 14], [16, 7, 7]], device='cuda', dtype=torch.int32)
+        >>> embeddings = rot_pos_emb(inv_freq, grids, spatial_merge_size=2)
+        >>> print(embeddings.shape)  # [computed_tokens, 128]
+
+    Note:
+        The function automatically dispatches to int32 or int64 implementations
+        based on the dtype of grid_thw. Output is always float32.
+    """
+    return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size)
+
+
+def get_window_index(
+    grid_thw: torch.Tensor,
+    spatial_merge_size: int,
+    vit_merger_window_size: int,
+    patch_size: int,
+    spatial_merge_unit: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate window attention indices for Vision Transformer architectures.
+
+    Computes window-based attention indices for hierarchical processing of vision
+    tokens, enabling efficient sliding window attention patterns in ViT models.
+
+    Args:
+        grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W)
+                                Must be int32 dtype on CUDA device
+        spatial_merge_size (int): Spatial dimension merge factor
+        vit_merger_window_size (int): Size of attention windows for ViT processing
+        patch_size (int): Size of vision patches in pixels
+        spatial_merge_unit (int): Unit size for spatial merging operations
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - window_indices: Flattened window indices of shape [total_elements]
+            - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1]
+
+    Example:
+        >>> grids = torch.tensor([[1, 14, 14]], device='cuda', dtype=torch.int32)
+        >>> indices, seqlens = get_window_index(
+        ...     grids, spatial_merge_size=2, vit_merger_window_size=7,
+        ...     patch_size=16, spatial_merge_unit=4
+        ... )
+
+    Note:
+        Returns empty tensors if input grid is empty or no valid windows can be formed.
+        The cu_window_seqlens tensor enables efficient batched attention computation.
+    """
+    return backend.get_window_index(
+        grid_thw,
+        spatial_merge_size,
+        vit_merger_window_size,
+        patch_size,
+        spatial_merge_unit,
+    )

wall_x/fusions/ops.py

@@ -495,3 +495,248 @@ class MultimodalRoPE(torch.autograd.Function):
 
 def multimodal_rope(q, k, cos, sin, mrope_section):
     return MultimodalRoPE.apply(q, k, cos, sin, mrope_section)
+
+
+################################################################################################
+##
+## RoPE Index 3D
+##
+################################################################################################
+
+
+def get_rope_index(
+    input_ids: torch.Tensor,
+    spatial_merge_size: int,
+    image_token_id: int,
+    video_token_id: int,
+    vision_start_token_id: int,
+    tokens_per_second: float,
+    image_grid_thw: torch.Tensor = None,
+    video_grid_thw: torch.Tensor = None,
+    second_per_grid_ts: torch.Tensor = None,
+    attention_mask: torch.Tensor = None,
+):
+    """
+    Generate 3D RoPE position indices for multimodal transformer inputs.
+
+    Computes position indices for text, image, and video tokens to enable proper
+    spatial-temporal position encoding in multimodal transformers with RoPE.
+
+    Args:
+        input_ids (torch.Tensor): Input token sequence of shape [batch_size, seq_len]
+                                 Must be LongTensor on CUDA device
+        spatial_merge_size (int): Spatial merge size for patch grouping (must be positive)
+        image_token_id (int): Token ID representing image patches
+        video_token_id (int): Token ID representing video frames
+        vision_start_token_id (int): Token ID marking start of vision sequences
+        tokens_per_second (float): Temporal scaling factor for video sequences (must be positive)
+        image_grid_thw (torch.Tensor, optional): Image grid dimensions of shape [num_images, 3] (T, H, W)
+        video_grid_thw (torch.Tensor, optional): Video grid dimensions of shape [num_videos, 3] (T, H, W)
+        second_per_grid_ts (torch.Tensor, optional): Video time intervals of shape [num_videos]
+        attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len]
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - position_ids: 3D position indices of shape [3, batch_size, seq_len]
+            - mrope_position_deltas: mRoPE position deltas of shape [batch_size, 1]
+
+    Raises:
+        TypeError: If input_ids is not a torch.Tensor
+        ValueError: If input dimensions are incorrect or tensors not on CUDA
+    """
+    # Input validation
+    if not isinstance(input_ids, torch.Tensor):
+        raise TypeError("input_ids must be a torch.Tensor")
+
+    if input_ids.dim() != 2:
+        raise ValueError("input_ids must be 2D tensor (batch_size, seq_len)")
+
+    if not input_ids.is_cuda:
+        raise ValueError("input_ids must be on CUDA device")
+
+    # Parameter validation
+    if not isinstance(spatial_merge_size, int) or spatial_merge_size <= 0:
+        raise ValueError(
+            f"spatial_merge_size must be positive integer, got {spatial_merge_size}"
+        )
+
+    if not isinstance(tokens_per_second, (int, float)) or tokens_per_second <= 0:
+        raise ValueError(
+            f"tokens_per_second must be positive number, got {tokens_per_second}"
+        )
+
+    return backend.get_rope_index(
+        input_ids,
+        image_grid_thw,
+        video_grid_thw,
+        second_per_grid_ts,
+        attention_mask,
+        spatial_merge_size,
+        image_token_id,
+        video_token_id,
+        vision_start_token_id,
+        float(tokens_per_second),
+    )
+
+
+################################################################################################
+##
+## Fused Rotary Position Embedding
+##
+################################################################################################
+
+
+def rot_pos_emb(
+    inv_freq: torch.Tensor,
+    grid_thw: torch.Tensor,
+    spatial_merge_size: int,
+) -> torch.Tensor:
+    """
+    Compute fused rotary position embeddings using optimized CUDA kernel.
+
+    This function fuses all rotary position embedding computations into a single
+    CUDA kernel for improved performance with spatial-temporal grids.
+
+    Args:
+        inv_freq (torch.Tensor): Inverse frequencies tensor of shape [dim/2]
+                                Contains precomputed 1.0 / (theta ** (torch.arange(0, dim, 2) / dim))
+                                Must be float32 on CUDA device
+        grid_thw (torch.Tensor): Grid dimensions tensor of shape [num_grids, 3]
+                                Each row contains (T, H, W) for temporal, height, width dimensions
+                                Supports int32 or int64 on CUDA device
+        spatial_merge_size (int): Spatial merge size for token grouping (must be positive)
+
+    Returns:
+        torch.Tensor: Rotary position embeddings of shape [total_tokens, dim]
+                     where dim = 2 * len(inv_freq)
+                     First half contains h_pos frequencies, second half contains w_pos frequencies
+
+    Raises:
+        TypeError: If inputs are not torch.Tensor or spatial_merge_size not int
+        ValueError: If tensor dimensions incorrect, not on CUDA, or devices mismatch
+        RuntimeError: If CUDA kernel execution fails
+    """
+    # Type checking
+    if not isinstance(inv_freq, torch.Tensor):
+        raise TypeError(f"inv_freq must be a torch.Tensor, got {type(inv_freq)}")
+
+    if not isinstance(grid_thw, torch.Tensor):
+        raise TypeError(f"grid_thw must be a torch.Tensor, got {type(grid_thw)}")
+
+    # Dimension checking
+    if inv_freq.dim() != 1:
+        raise ValueError(
+            f"inv_freq must be 1-dimensional, got {inv_freq.dim()}D tensor"
+        )
+
+    if grid_thw.dim() != 2:
+        raise ValueError(
+            f"grid_thw must be 2-dimensional, got {grid_thw.dim()}D tensor"
+        )
+
+    if grid_thw.size(1) != 3:
+        raise ValueError(
+            f"grid_thw must have shape [num_grids, 3], got shape {list(grid_thw.shape)}"
+        )
+
+    # Device checking
+    if not inv_freq.is_cuda:
+        raise ValueError("inv_freq must be on CUDA device")
+
+    if not grid_thw.is_cuda:
+        raise ValueError("grid_thw must be on CUDA device")
+
+    # Ensure both tensors are on the same device
+    if inv_freq.device != grid_thw.device:
+        raise ValueError(
+            f"inv_freq and grid_thw must be on the same device, "
+            f"got {inv_freq.device} and {grid_thw.device}"
+        )
+
+    # Parameter validation
+    if not isinstance(spatial_merge_size, int):
+        raise TypeError(
+            f"spatial_merge_size must be an integer, got {type(spatial_merge_size)}"
+        )
+
+    if spatial_merge_size <= 0:
+        raise ValueError(
+            f"spatial_merge_size must be positive, got {spatial_merge_size}"
+        )
+
+    # Ensure inv_freq is float32 (the kernel expects float)
+    if inv_freq.dtype != torch.float32:
+        inv_freq = inv_freq.to(torch.float32)
+
+    # Call the CUDA backend
+    try:
+        return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size)
+    except RuntimeError as e:
+        raise RuntimeError(f"CUDA kernel execution failed: {str(e)}")
+
+
+################################################################################################
+##
+## Fused Window Index Generation
+##
+################################################################################################
+
+
+def get_window_index(
+    grid_thw: torch.Tensor,
+    window_size: int,
+    spatial_merge_size: int,
+    patch_size: int,
+    spatial_merge_unit: int = 1,
+):
+    """
+    Generate window attention indices for Vision Transformer architectures.
+
+    Computes window-based attention indices for hierarchical processing of vision
+    tokens, enabling efficient sliding window attention patterns in ViT models.
+
+    Args:
+        grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W)
+                                Must be or will be converted to int32 on CUDA device
+        window_size (int): Window size for attention computation
+        spatial_merge_size (int): Spatial merge size for patch grouping
+        patch_size (int): Size of vision patches in pixels
+        spatial_merge_unit (int, optional): Spatial merging unit size. Defaults to 1.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - window_index: Window indices tensor of shape [total_elements]
+            - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1]
+
+    Raises:
+        AssertionError: If grid_thw dimensions are incorrect
+
+    Note:
+        Returns empty tensors if input grid is empty or no valid windows can be formed.
+        The function automatically converts input to CUDA int32 if needed.
+    """
+    # Input validation
+    assert (
+        grid_thw.dim() == 2 and grid_thw.size(1) == 3
+    ), f"grid_thw must have shape (num_grids, 3), got {grid_thw.shape}"
+
+    # Ensure input is on CUDA and int32 type
+    if not grid_thw.is_cuda:
+        grid_thw = grid_thw.cuda()
+
+    if grid_thw.dtype != torch.int32:
+        grid_thw = grid_thw.to(torch.int32)
+
+    # Calculate vit_merger_window_size
+    vit_merger_window_size = window_size // spatial_merge_size // patch_size
+
+    # Call CUDA backend
+    window_index, cu_window_seqlens = backend.get_window_index(
+        grid_thw,
+        spatial_merge_size,
+        vit_merger_window_size,
+        patch_size,
+        spatial_merge_unit,
+    )
+
+    return window_index, cu_window_seqlens

wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py

@@ -560,15 +560,17 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel):
             `torch.Tensor`: hidden_states.
         """
         hidden_states = self.patch_embed(hidden_states)
-        rotary_pos_emb = self.rot_pos_emb(grid_thw)
-        window_index, cu_window_seqlens = self.get_window_index(grid_thw)
-        window_index = window_index.to(hidden_states.device)
-        cu_window_seqlens = torch.tensor(
-            cu_window_seqlens,
-            device=hidden_states.device,
-            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+        rotary_pos_emb = ops.rot_pos_emb(
+            self.rotary_pos_emb.inv_freq, grid_thw, self.spatial_merge_size
+        )
+
+        window_index, cu_window_seqlens = ops.get_window_index(
+            grid_thw=grid_thw,
+            window_size=self.window_size,
+            spatial_merge_size=self.spatial_merge_size,
+            patch_size=self.patch_size,
+            spatial_merge_unit=self.spatial_merge_unit,
         )
-        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
 
         seq_len, _ = hidden_states.size()
         hidden_states = hidden_states.reshape(

wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py

@@ -1244,12 +1244,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                 or self.rope_deltas is None
                 or (past_key_values is None or past_key_values.get_seq_length() == 0)
             ):
-                position_ids, rope_deltas = self.get_rope_index(
-                    input_ids,
-                    image_grid_thw,
-                    video_grid_thw,
-                    second_per_grid_ts,
-                    attention_mask,
+                position_ids, rope_deltas = ops.get_rope_index(
+                    input_ids=input_ids,
+                    image_grid_thw=image_grid_thw,
+                    video_grid_thw=video_grid_thw,
+                    second_per_grid_ts=second_per_grid_ts,
+                    attention_mask=attention_mask,
+                    spatial_merge_size=self.config.vision_config.spatial_merge_size,
+                    image_token_id=self.config.image_token_id,
+                    video_token_id=self.config.video_token_id,
+                    vision_start_token_id=self.config.vision_start_token_id,
+                    tokens_per_second=self.config.vision_config.tokens_per_second,
                 )
                 self.rope_deltas = rope_deltas
             # Use previously calculated rope deltas to get correct position IDs
@@ -1720,12 +1725,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                 or self.rope_deltas is None
                 or (past_key_values is None or past_key_values.get_seq_length() == 0)
             ):
-                position_ids, rope_deltas = self.get_rope_index(
-                    input_ids,
-                    image_grid_thw,
-                    video_grid_thw,
-                    second_per_grid_ts,
-                    attention_mask,
+                position_ids, rope_deltas = ops.get_rope_index(
+                    input_ids=input_ids,
+                    image_grid_thw=image_grid_thw,
+                    video_grid_thw=video_grid_thw,
+                    second_per_grid_ts=second_per_grid_ts,
+                    attention_mask=attention_mask,
+                    spatial_merge_size=self.config.vision_config.spatial_merge_size,
+                    image_token_id=self.config.image_token_id,
+                    video_token_id=self.config.video_token_id,
+                    vision_start_token_id=self.config.vision_start_token_id,
+                    tokens_per_second=self.config.vision_config.tokens_per_second,
                 )
                 self.rope_deltas = rope_deltas
             # Use previously calculated rope deltas to get correct position IDs
@@ -1882,6 +1892,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
             )
             dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype)
 
+            # Calculate token distribution across MoE expert groups
+            group_size = torch.zeros(
+                self.config.num_experts, dtype=torch.long, device="cpu"
+            )
+            for i in range(self.config.num_experts):
+                group_size[i] = (moe_token_types == i).sum()
+
+            # Calculate start and end indices for each expert group
+            start_indices = torch.cumsum(group_size, dim=0) - group_size
+            end_indices = torch.cumsum(group_size, dim=0)
+
             def step(timestep, noisy_action):
                 """
                 Single denoising step for diffusion process.
@@ -1915,6 +1936,8 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                     past_key_values=past_key_values,
                     inputs_embeds=temp_inputs_embeds,
                     moe_token_types=moe_token_types,
+                    start_indices=start_indices,
+                    end_indices=end_indices,
                     use_cache=True,
                     output_attentions=False,
                     output_hidden_states=False,

wall_x/trainer/qwen_vl_act_trainer.py

@@ -5,11 +5,13 @@ import torch
 import random
 import numpy as np
 import torch.nn as nn
+import torch.distributed as dist
 
 from tqdm import tqdm
 from functools import wraps
 from datetime import datetime
 from torch.optim import AdamW
+from torch.distributed.tensor import distribute_tensor
 from accelerate import Accelerator
 from safetensors.torch import load_file
 from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
@@ -769,7 +771,9 @@ class QwenVlAct_Trainer:
         """
         checkpoint_path = self.config["resume"]["ckpt"]
 
-        if self.config.get("resume", {}).get("load_ckpt_only", False):
+        if self.config.get("FSDP2", False):
+            self._load_fsdp_state_dict_with_distribute_tensor()
+        elif self.config.get("resume", {}).get("load_ckpt_only", False):
             # Load only model weights
             ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors"
             state_dict = load_file(ckpt_path, device="cpu")
@@ -788,6 +792,48 @@ class QwenVlAct_Trainer:
 
         self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")
 
+    def _load_fsdp_state_dict_with_distribute_tensor(self):
+
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        full_sd = load_file(
+            self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu"
+        )
+        meta_sharded_sd = self.model.state_dict()
+        sharded_sd = {}
+
+        def find_matching_key(target_key, available_keys):
+            if target_key in available_keys:
+                return target_key
+            prefixed_key = f"_orig_mod.{target_key}"
+            if prefixed_key in available_keys:
+                return prefixed_key
+            if target_key.startswith("_orig_mod."):
+                unprefixed_key = target_key[len("_orig_mod.") :]
+                if unprefixed_key in available_keys:
+                    return unprefixed_key
+            return None
+
+        for param_name, full_tensor in full_sd.items():
+            matching_key = find_matching_key(param_name, meta_sharded_sd.keys())
+            if matching_key is None:
+                if rank == 0:
+                    print(
+                        f"[Rank {rank}] Warning: Parameter not found:",
+                        param_name,
+                        flush=True,
+                    )
+                continue
+            sharded_meta_param = meta_sharded_sd[matching_key]
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                sharded_meta_param.device_mesh,
+                sharded_meta_param.placements,
+            )
+            sharded_sd[matching_key] = nn.Parameter(sharded_tensor)
+
+        self.model.load_state_dict(sharded_sd, assign=True, strict=False)
+
     def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask):
         """
         Log detailed L1 loss metrics by degrees of freedom.