NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Is forward_step_func expected to be skipped in combined_1f1b_schedule_for_interleaved_pipelining?

Open
#1,830 0 comments 1 reaction 0 assignees View on GitHub
question
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

Hi,

While reading the code in combined_1f1b.py, especially the function combined_1f1b_schedule_for_interleaved_pipelining, I noticed that it seems to bypass the user-provided forward_step_func.

Normally, forward_step_func is where we place hooks such as straggler timers. For example:

```
def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = False):
args = get_args()
timers = get_timers()

global stimer
with stimer(bdata=True):
vp_stage = get_attr_wrapped_model(model, "vp_stage")
tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator, vp_stage)

with stimer:
if return_schedule_plan:
assert args.overlap_moe_expert_parallel_comm
schedule_plan = model.build_schedule_plan(
tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
)
return schedule_plan, partial(loss_func, loss_mask, model=model)
else:
output_tensor = model(
tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
)

return output_tensor, partial(loss_func, loss_mask, model=model)
```

In the regular 1F1B schedule this function is invoked, so timers are triggered correctly.
But in combined_1f1b_schedule_for_interleaved_pipelining, from the code path it looks like the schedule plan and losses are constructed directly, without going through forward_step_func.

Questions:

Is skipping forward_step_func in the combined interleaved schedule intentional?

If so, what is the recommended way to insert timing hooks (e.g., for straggler mitigation) in this schedule?

If not intentional, should this function call be added to maintain consistency with the other schedules?

I haven’t run into this through execution yet — just noticed it while reading the code. Please let me know if I misunderstood the intended design.

Thanks!

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.