Helper method to get the model chunk ID given the iteration number.
(microbatch_id, forward)
| 211 | num_microbatches_remaining = num_microbatches - num_warmup_microbatches |
| 212 | |
| 213 | def get_model_chunk_id(microbatch_id, forward): |
| 214 | """Helper method to get the model chunk ID given the iteration number.""" |
| 215 | microbatch_id_in_group = microbatch_id % ( |
| 216 | pipeline_parallel_size * num_model_chunks |
| 217 | ) |
| 218 | model_chunk_id = microbatch_id_in_group // pipeline_parallel_size |
| 219 | if not forward: |
| 220 | model_chunk_id = num_model_chunks - model_chunk_id - 1 |
| 221 | return model_chunk_id |
| 222 | |
| 223 | def forward_step_helper(microbatch_id): |
| 224 | """Helper method to run forward step with model split into chunks |
no outgoing calls
no test coverage detected