actaug GR00T train+eval bundle: code (save-at/stop-at), eval stack (porting fixes), subset builder, protocol docs
38dc077 verified | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| #!/bin/bash | |
| echo "Important Notes:" | |
| echo "1: The max batch of engine size is set to 8 in the reference case. " | |
| echo "2: The MIN_LEN/OPT_LEN/MAX_LEN for LLM, DiT, VLLN-VLSelfAttention models is set to 80/283/300." | |
| echo "If your inference batch size exceeds 8 or the MIN_LEN/OPT_LEN/MAX_LEN for LLM, DiT, VLLN-VLSelfAttention not fit your use case, please set it to your actual batch size and length variables." | |
| export PATH=/usr/src/tensorrt/bin:$PATH | |
| # Define length variables | |
| MIN_LEN=80 | |
| OPT_LEN=283 | |
| MAX_LEN=300 | |
| if [ -e /usr/src/tensorrt/bin/trtexec ]; then | |
| echo "The file /usr/src/tensorrt/bin/trtexec exists." | |
| else | |
| echo "The file /usr/src/tensorrt/bin/trtexec does not exist. Please install tensorrt" | |
| fi | |
| mkdir -p gr00t_engine | |
| # VLLN-VLSelfAttention | |
| echo "------------Building vlln_vl_self_attention Model--------------------" | |
| trtexec --useCudaGraph --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/action_head/vlln_vl_self_attention.onnx --saveEngine=gr00t_engine/vlln_vl_self_attention.engine --minShapes=backbone_features:1x${MIN_LEN}x2048 --optShapes=backbone_features:1x${OPT_LEN}x2048 --maxShapes=backbone_features:8x${MAX_LEN}x2048 > gr00t_engine/vlln_vl_self_attention.log 2>&1 | |
| # DiT Model | |
| echo "------------Building DiT Model--------------------" | |
| trtexec --useCudaGraph --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/action_head/DiT.onnx --saveEngine=gr00t_engine/DiT.engine --minShapes=sa_embs:1x17x1536,vl_embs:1x${MIN_LEN}x2048,timesteps_tensor:1 --optShapes=sa_embs:1x17x1536,vl_embs:1x${OPT_LEN}x2048,timesteps_tensor:1 --maxShapes=sa_embs:8x17x1536,vl_embs:8x${MAX_LEN}x2048,timesteps_tensor:8 > gr00t_engine/build_DiT.log 2>&1 | |
| # State Encoder | |
| echo "------------Building State Encoder--------------------" | |
| trtexec --useCudaGraph --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/action_head/state_encoder.onnx --saveEngine=gr00t_engine/state_encoder.engine --minShapes=state:1x1x64,embodiment_id:1 --optShapes=state:1x1x64,embodiment_id:1 --maxShapes=state:8x1x64,embodiment_id:8 > gr00t_engine/build_state_encoder.log 2>&1 | |
| # Action Encoder | |
| echo "------------Building Action Encoder--------------------" | |
| trtexec --useCudaGraph --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/action_head/action_encoder.onnx --saveEngine=gr00t_engine/action_encoder.engine --minShapes=actions:1x16x32,timesteps_tensor:1,embodiment_id:1 --optShapes=actions:1x16x32,timesteps_tensor:1,embodiment_id:1 --maxShapes=actions:8x16x32,timesteps_tensor:8,embodiment_id:8 > gr00t_engine/build_action_encoder.log 2>&1 | |
| # Action Decoder | |
| echo "------------Building Action Decoder--------------------" | |
| trtexec --useCudaGraph --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/action_head/action_decoder.onnx --saveEngine=gr00t_engine/action_decoder.engine --minShapes=model_output:1x17x1024,embodiment_id:1 --optShapes=model_output:1x17x1024,embodiment_id:1 --maxShapes=model_output:8x17x1024,embodiment_id:8 > gr00t_engine/build_action_decoder.log 2>&1 | |
| # VLM-ViT | |
| echo "------------Building VLM-ViT--------------------" | |
| trtexec --useCudaGraph --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/eagle2/vit.onnx --saveEngine=gr00t_engine/vit.engine --minShapes=pixel_values:1x3x224x224,position_ids:1x256 --optShapes=pixel_values:1x3x224x224,position_ids:1x256 --maxShapes=pixel_values:8x3x224x224,position_ids:8x256 > gr00t_engine/vit.log 2>&1 | |
| # VLM-LLM | |
| echo "------------Building VLM-LLM--------------------" | |
| trtexec --verbose --stronglyTyped --separateProfileRun --noDataTransfers --onnx=gr00t_onnx/eagle2/llm.onnx --saveEngine=gr00t_engine/llm.engine --minShapes=input_ids:1x${MIN_LEN},vit_embeds:1x256x1152,attention_mask:1x${MIN_LEN} --optShapes=input_ids:1x${OPT_LEN},vit_embeds:1x256x1152,attention_mask:1x${OPT_LEN} --maxShapes=input_ids:8x${MAX_LEN},vit_embeds:8x256x1152,attention_mask:8x${MAX_LEN} > gr00t_engine/llm.log 2>&1 | |