feat: speed up training 20260828 - #378
Draft
SakodaShintaro wants to merge 2 commits into
Draft
SakodaShintaro wants to merge 2 commits into
SakodaShintaro wants to merge 2 commits into
Conversation
The per-element MLP-Mixer blocks in EgoEncoder / NeighborEncoder /
LaneEncoder / LineEncoder ran at a hardcoded 128 channels and 64 tokens
for `encoder_mixer_depth` = 6 layers. They are applied to every element
(~560 per sample) at every point (20-31), and are LayerNorm / GELU /
permute heavy, so their activations -- not the 256-dim hidden state --
dominated encoder time and memory.
Make both mixer sizes configurable alongside the existing depth and
default them to 32 / 32 / 1. The sizes are threaded from ModelConfig
through Encoder into the four sequence encoders; the encoders without
mixer blocks (StaticEncoder / GoalPoseEncoder / FloatsEncoder) keep
their 128-dim projection, since shrinking those buys no speed.
Encoder forward+backward, bf16 autocast, RTX 4080, 20 valid neighbours
(the dataset median):
batch 128/64/6 32/32/1
8 55.5 ms 2.85 GiB 33.5 ms 0.47 GiB
32 223.0 ms 10.90 GiB 54.3 ms 1.67 GiB
64 out of memory 81.0 ms 3.25 GiB
Encoder parameters drop from 7.42M to 5.66M. The training gather path
and the eval / ONNX mask path still agree to 2.4e-06, and ONNX export
succeeds at both sizes.
Checkpoints trained before this change need
`encoder_mixer_hidden_dim: 128`, `encoder_mixer_token_dim: 64` and
`encoder_mixer_depth: 6` added to their args.json to be re-exported or
validated.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Shintaro Sakoda <shintaro.sakoda@tier4.jp>
…o 4096 train_epoch ran the forward in fp32 with no way to switch. Wrap the loss computation in `torch.autocast(bfloat16)` behind a new `use_amp` flag, defaulted to on: bf16 keeps the fp32 exponent range, so no GradScaler is needed, and the backward stays in fp32 through autocast's own casting. `batch_size` is already on the command line; raise its default from 512 to 4096. It is the batch across all GPUs (train.py divides by the world size), and the smaller encoder mixer from the previous commit cut encoder activations by ~6.5x, so the previous default no longer reflects what fits. Pass `--use_amp False` to get the old fp32 behaviour back. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Shintaro Sakoda <shintaro.sakoda@tier4.jp>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
use_ampflag and raises thebatch_sizedefault from 512 to 4096