[QUESTION]Cudagraph support for Transformer_Engine
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Your question**
I have a question regarding cudagraph functionality in the codebase. I could see that in the code below that "--cuda-graph-impl local" option does not provide cudagraph support during (SFT) training. I also tried "--cuda-graph-impl transformer_engine" option, it does not seem to support it too. Although, it is not clear from the code if cudagraph is supported or not for transfomer_engine option. Please confirm if cudagraph is supported for transformer_engine option.
Also, are their any implications of --use-dynamic-batch-size --max-tokens-per-gpu 4096 options on cudagraph creation.
Appreciate it. Thanks.
++++++++++
def _should_call_local_cudagraph(self, *args, **kwargs):
"""
Check if we should call the local cudagraph path.
"""
# Training and validation mode CUDA graphs
if hasattr(self, 'cudagraph_manager') and kwargs.get('inference_context') is None:
return True
# Inference mode. CUDA graphs are used in the decode phase only, when attn mask is None
elif not self.training and (
hasattr(self, 'cudagraph_manager')
and kwargs['attention_mask'] is None
and (
(kwargs.get('inference_context') is not None)
or (kwargs.get('inference_params') is not None)
)
and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope
):
if kwargs['inference_context'].is_static_batching():
using_cuda_graph = kwargs['inference_context'].is_decode_only()
else:
# it can happen that non-decode steps have a token count greater than the max
# supported cuda graph token count. In that case this flag will be set to
# False by initialize_attention, and we should not use cuda graphs.
using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step()
if using_cuda_graph:
return True
return False
Contributor guide
Assessment
This issue has not been assessed yet.