Diff Coverage

Diff: origin/master...HEAD, staged and unstaged changes

Source File Diff Coverage (%) Missing Lines
hyper_parallel/components/losses/_vocab_parallel_cross_entropy.py 100%  
hyper_parallel/components/losses/dispatcher.py 100%  
hyper_parallel/components/losses/mtp.py 0.0% 43
hyper_parallel/core/activation_memory/api.py 100%  
hyper_parallel/core/pipeline_parallel/_microbatch.py 91.7% 59,73
hyper_parallel/core/pipeline_parallel/_stage.py 66.7% 247
hyper_parallel/components/losses/mtp.py
39
40
41
42
43
44
45
46
47
    Returns:
        Summed MTP loss over all depths.
    """
    total_mtp_loss = torch.tensor(0.0, device=labels.device, dtype=torch.float32)
    for logits in mtp_per_depth_logits:
        logits_shifted = logits[..., :-1, :].contiguous()
        labels_shifted = labels[..., 1:].contiguous()
        depth_loss = loss_fn(
            logits_shifted.view(-1, logits_shifted.size(-1)),
hyper_parallel/core/pipeline_parallel/_microbatch.py
55
56
57
58
59
60
61
62
63
            cur_arg_batch_dim = 0
            if self.args_batch_dim and self.args_batch_dim[arg_idx] is not None:
                cur_arg_batch_dim = self.args_batch_dim[arg_idx].batch_dim
            if isinstance(cur_arg, hyper_parallel.DTensor):
                micro_arg = self.split_inputs_with_custom_shard(cur_arg, cur_arg_batch_dim, micro_idx)
            else:
                micro_arg = self.split_inputs(cur_arg, cur_arg_batch_dim, micro_idx)
            micro_args.append(micro_arg)
        return micro_args

69
70
71
72
73
74
75
76
77
            cur_kwarg_batch_dim = 0
            if self.kwargs_batch_dim is not None:
                cur_kwarg_batch_dim = self.kwargs_batch_dim[key].batch_dim
            if isinstance(cur_kwarg, hyper_parallel.DTensor):
                micro_kwarg = self.split_inputs_with_custom_shard(cur_kwarg, cur_kwarg_batch_dim, micro_idx)
            else:
                micro_kwarg = self.split_inputs(cur_kwarg, cur_kwarg_batch_dim, micro_idx)
            micro_kwargs[key] = micro_kwarg
        return micro_kwargs
hyper_parallel/core/pipeline_parallel/_stage.py
243
244
245
246
247
248
249
250
251
        """Get last stage sens"""
        p_sens = None
        if isinstance(last_stage_outputs, (list, tuple)):
            p_sens = []
            for out_i in last_stage_outputs:
                if isinstance(out_i, hyper_parallel.DTensor):
                    repeat_num = out_i.layout.repeat_num()
                    sens_i = torch.full_like(out_i.to_local(), 1.0 / repeat_num)
                else: