fix: Mooncake RDMA memory for aux_layers > 3#140
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 52afe0c571
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| num_target_layers = getattr(draft_cfg, "num_target_layers", None) | ||
| if num_target_layers: | ||
| return int(num_target_layers) | ||
| aux_layers = getattr(args, "inference_aux_hidden_states_layers", None) |
There was a problem hiding this comment.
When users set inference.aux_hidden_states_layers in the YAML, config_to_flat_args exposes it as args.aux_hidden_states_layers, not args.inference_aux_hidden_states_layers. In configs whose draft config does not carry target_layer_ids/num_target_layers (for example standard Eagle3 configs with more than three explicit aux layers), this branch misses the user setting and falls back to 3, so the exported MOONCAKE_NUM_AUX_LAYERS and auto host/GPU buffer sizes remain too small and the same >3-layer Mooncake allocation failure can still occur.
Useful? React with 👍 / 👎.
Signed-off-by: Doğaç Eldenk <dogacel@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b8f09a49ca
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| target_layer_ids = getattr(draft_cfg, "target_layer_ids", None) | ||
| if target_layer_ids: | ||
| return len(target_layer_ids) | ||
| num_target_layers = getattr(draft_cfg, "num_target_layers", None) | ||
| if num_target_layers: | ||
| return int(num_target_layers) |
There was a problem hiding this comment.
In DFlash/DSpark runs where inference.aux_hidden_states_layers is set to override the draft config (for example, capturing more layers than the config's target_layer_ids), train_entry leaves that explicit list on args and the inference engines use it for capture, but this branch returns the draft config's length before looking at args.aux_hidden_states_layers. That sizes MOONCAKE_HOST_BUFFER_SIZE/GPU buffers and exports MOONCAKE_NUM_AUX_LAYERS for the wrong layer count, so the allocation failure this change is meant to fix can still occur; check the explicit aux layer list before falling back to the draft config.
Useful? React with 👍 / 👎.
2 participants
num_aux_layersto be 3, which causes OOM mid-training. Fixed it by parsing draft config properly.Especially warm start helps us start directly from pre-trained hugging-face checkpoints, including DFlash checkpoint warm starting a DSpark model.