loading pretrained model with bfloat16 during fine-tuning on 24GB GPU

Describe the bug

I noticed that in gr00t/model/gr00t_n1d7/setup.py (around line 82), the pretrained model is loaded with AutoModel.from_pretrained without specifying bfloat16 dtype.

Since I only have a 24GB GPU, I modified the code and added:

dtype=torch.bfloat16

to load the model in bfloat16 format before fine-tuning, which allows me to fit the training process into the available VRAM.

I would like to confirm:

1. Is loading the pretrained model with bfloat16 recommended for fine-tuning?
2. Does using bfloat16 during model loading have any impact on training stability or final performance compared with the default loading method?
3. Are there any recommended settings for fine-tuning GR00T on GPUs with limited VRAM?

Additionally, after making this change, my training crashed twice during the fine-tuning process. I am not sure whether this is related to using bfloat16.

Could you please advise if bfloat16 loading is supported for fine-tuning and whether there are any known issues?

Thanks!

Versions

Python 3.12.13 contourpy==1.3.3 controlnet_aux==0.0.7 coverage==7.4.4 croniter==6.0.0 cryptography==46.0.7 cssselect==1.3.0 cssselect2==0.8.0 cssutils==2.11.1 ctranslate2==4.6.2 cuda-bindings==12.9.4 cuda-pathfinder==1.5.5 cuda-python==13.0.2 cuda-toolkit==13.0.2 -e git+https://github.com/FindDefinition/cumm@4c77b38d1ab57d5d1c157adddf67dad93f3a446b#egg=cumm -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab&subdirectory=source/isaaclab

Editable install with no version control (isaaclab3_demo==0.1.0)

-e /home/zzg/workspace/pycharm/isaaclab3_demo/source/isaaclab3_demo

Editable install with no version control (isaaclab3_loco_go2==0.1.0)

-e /home/zzg/workspace/pycharm/isaaclab3_loco_go2/source/isaaclab3_loco_go2 -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_assets&subdirectory=source/isaaclab_assets -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_contrib&subdirectory=source/isaaclab_contrib -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_experimental&subdirectory=source/isaaclab_experimental -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_mimic&subdirectory=source/isaaclab_mimic -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_newton&subdirectory=source/isaaclab_newton -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_ov&subdirectory=source/isaaclab_ov -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_ovphysx&subdirectory=source/isaaclab_ovphysx -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_physx&subdirectory=source/isaaclab_physx -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_ppisp&subdirectory=source/isaaclab_ppisp -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_rl&subdirectory=source/isaaclab_rl -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_tasks&subdirectory=source/isaaclab_tasks -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_tasks_experimental&subdirectory=source/isaaclab_tasks_experimental -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_teleop&subdirectory=source/isaaclab_teleop -e git+https://github.com/isaac-sim/IsaacLab.git@0954ba7e2445f5b0902e3f4b948ebcd5d48b8ee3#egg=isaaclab_visualizers&subdirectory=source/isaaclab_visualizers isaacsim==6.0.0.1 isaacsim-app==6.0.0.1 isaacsim-asset==6.0.0.1 isaacsim-benchmark==6.0.0.1 isaacsim-code-editor==6.0.0.1 isaacsim-core==6.0.0.1 isaacsim-cortex==6.0.0.1 isaacsim-example==6.0.0.1 isaacsim-extscache-kit==6.0.0.1 isaacsim-extscache-kit-sdk==6.0.0.1 isaacsim-extscache-physics==6.0.0.1 isaacsim-gui==6.0.0.1 isaacsim-kernel==6.0.0.1 isaacsim-replicator==6.0.0.1 isaacsim-rl==6.0.0.1 isaacsim-robot==6.0.0.1 isaacsim-robot-motion==6.0.0.1 isaacsim-robot-setup==6.0.0.1 isaacsim-ros1==6.0.0.1 isaacsim-ros2==6.0.0.1 isaacsim-sensor==6.0.0.1 isaacsim-storage==6.0.0.1 isaacsim-template==6.0.0.1 isaacsim-test==6.0.0.1 isaacsim-utils==6.0.0.1 isaacteleop==1.3.131 newton-usd-schemas==0.2.0 nibabel==5.3.3 ninja==1.13.0 nlopt==2.10.0 nltk==3.9.1 nnAudio==0.3.3 nodeenv==1.9.1 notebook==6.5.7 notebook_shim==0.2.4 num2words==0.5.14 numba==0.61.2 numcodecs==0.16.5 numexpr==2.14.1 numpy==2.3.1 numpy-quaternion==2024.0.13 numpy-stl==4.0.0 numpydantic==1.8.1 nuscenes-devkit==1.2.0 nvidia-arch==7.1.0 nvidia-cublas==13.1.0.3 nvidia-cublas-cu12==12.8.4.1 nvidia-cuda-cccl==13.0.85 nvidia-cuda-cupti==13.0.85 nvidia-cuda-cupti-cu12==12.8.90 nvidia-cuda-nvcc-cu12==12.9.86 nvidia-cuda-nvrtc==13.0.88 nvidia-cuda-nvrtc-cu12==12.8.93 nvidia-cuda-runtime==13.0.96 nvidia-cuda-runtime-cu12==12.8.90 nvidia-cudnn-cu12==9.10.2.21 nvidia-cudnn-cu13==9.19.0.56 nvidia-cudnn-frontend==1.15.0 nvidia-cufft==12.0.0.61 nvidia-cufft-cu12==11.3.3.83 nvidia-cufile==1.15.1.6 nvidia-cufile-cu12==1.13.1.3 nvidia-curand==10.4.0.35 nvidia-curand-cu12==10.3.9.90 -e git+https://github.com/NVlabs/curobo@ec2bfa9f9ea2673abe00cf911ef39353ddb79c29#egg=nvidia_curobo nvidia-cusolver==12.0.4.66 nvidia-cusolver-cu12==11.7.3.90 nvidia-cusparse==12.6.3.3 nvidia-cusparse-cu12==12.5.8.93 nvidia-cusparselt-cu12==0.7.1 nvidia-cusparselt-cu13==0.8.0 nvidia-cutlass-dsl==4.5.0.dev0 nvidia-cutlass-dsl-libs-base==4.5.0.dev0 nvidia-cutlass-dsl-libs-cu13==4.5.0.dev0 nvidia-ml-py==12.575.51 nvidia-nccl-cu12==2.27.5 nvidia-nccl-cu13==2.28.9 nvidia-nvjitlink==13.0.88 nvidia-nvjitlink-cu12==12.8.93 nvidia-nvshmem-cu12==3.4.5 nvidia-nvshmem-cu13==3.4.5 nvidia-nvtx==13.0.85 nvidia-nvtx-cu12==12.8.90 PyMySQL==1.1.1 PyNaCl==1.6.0 python-box==6.1.0 python-can==4.6.1 python-crfsuite==0.9.11 python-dateutil==2.9.0.post0 python-docx==1.2.0 python-dotenv==1.2.1 python-engineio==4.12.3 python-iso639==2025.2.18 python-jose==3.5.0 python-json-logger==4.0.0 python-magic==0.4.27 python-multipart==0.0.22 python-nightscout==1.0.0 python-oxmsg==0.0.2 python-pptx==1.0.2 python-slugify==8.0.4 python-socketio==5.13.0 python-socks==2.7.2 python-stdnum==1.20 python-utils==3.9.1 python-xlib==0.33 pytigre @ file:///home/zzg/workspace/pycharm/TIGRE pytokens==0.3.0 pytools==2025.2.5 pytorch-fid==0.3.0 pytorch-kinematics==0.7.6 pytorch-lightning==2.5.5 pytorch-msssim==1.0.0 pytorch-seed==0.2.0 pytorch-triton==3.5.0+git7416ffcb pytorch_optimizer==3.9.0 pytransform3d==3.15.0 pytube==15.0.0 pytz==2024.1 pyvers==0.2.2 pyviz_comms==3.0.6 PyWavelets==1.9.0 rerun-sdk==0.33.1 resampy==0.4.3 RestrictedPython==8.0 retrying==1.4.2 rf-clip==1.1 rfc3339-validator==0.1.4 rfc3986-validator==0.1.1 rfc3987-syntax==1.1.0 rhoban-cmeel-jsoncpp==1.9.4.9 rich==14.3.4 rich-rst==1.3.2 rich-toolkit==0.15.1 rignore==0.7.0 rl_games @ git+https://github.com/isaac-sim/rl_games.git@6b3534f29568158e9e29ec8bf83cc88fce5f0cae rlinf==0.2.0.dev2 sapien==3.0.3 schedulefree==1.4 scikit-image==0.25.2 scikit-learn==1.7.2 scipy==1.17.0 seaborn==0.13.2 SecretStorage==3.4.0 segment-anything==1.0 segmentation_models_pytorch==0.5.0 semantic-version==2.10.0 semver==3.0.4 Send2Trash==1.8.3 sensai-utils==1.6.0 sentence-transformers==5.1.1 sentencepiece==0.2.1 sentry-sdk==2.42.1 setproctitle==1.3.7 setuptools==81.0.0 setuptools-scm==9.2.2 sgl-kernel==0.3.15 sglang==0.5.3.post1 sgmllib3k==1.0.0 shap==0.51.0 shapely==2.0.7 SharedArray==3.2.4 shellingham==1.5.4 streamlit==1.51.0 StrEnum==0.4.15 strictyaml==1.7.3 stringzilla==4.2.1 structlog==24.4.0 submitit==1.5.3 supabase==2.27.2 supabase-auth==2.27.2 supabase-functions==2.27.2 supervision==0.26.1 surrealist==1.1.2 svg.path==7.0 svglib==1.6.0 sympy==1.14.0 tabulate @ file:///home/conda/feedstock_root/build_artifacts/tabulate_1733589744265/work tavily-python==0.7.13 teleop==0.1.5 tenacity==8.5.0 tensorboard==2.20.0 tensorboard-data-server==0.7.2 tensorboardX==2.6.4 tensordict==0.13.0 tensorly==0.9.0 tensorstore==0.1.74 termcolor==3.3.0 terminado==0.18.1 terminaltables==3.1.10 testpath==0.6.0 text-unidecode==1.3 textblob==0.19.0 textual==8.1.1 thop==0.1.1.post2209072238 threadpoolctl==3.6.0 tianshou==2.0.1 tidepool-data-science-project==0.0.3 tifffile==2025.10.16 tiktoken==0.12.0 timm==1.0.22 tinycss2==1.4.0 tinyobjloader==2.0.0rc13 tldextract==5.1.3 tokenizers==0.23.0rc0 toml==0.10.2 tomli==2.4.1 tomlkit==0.12.0 toolz==1.1.0 toppra==0.6.3 torch==2.10.0+cu128 torch-audiomentations==0.12.0 torch-complex==0.4.4 torch-einops-utils==0.0.30 torch-fidelity==0.3.0 torch-geometric==2.7.0 torch-stoi==0.2.3 torch_c_dlpack_ext==0.1.5 torch_cluster==1.6.3+pt210cu130 torch_memory_saver==0.0.9rc2 torch_pitch_shift==1.2.5 torch_scatter==2.1.2+pt210cu130 torch_sparse==0.6.18+pt210cu130 torch_spline_conv==1.2.2+pt210cu130 torchac==0.9.3 torchao==0.13.0 torchaudio==2.11.0 torchcodec==0.10.0+cu130 torchdiffeq==0.2.5 torchinfo==1.8.0 torchmetrics==1.8.2 torchreid==0.2.5 torchscale==0.3.0 torchsde==0.2.6 torchstat==0.0.7 torchsummary==1.5.1 torchvision==0.25.0+cu128 tornado==6.5.2 tqdm==4.67.1 tqdm-loggable==0.4.1 trackio==0.22.0 traitlets==5.14.3 trampoline==0.1.2 transformers==4.57.6 transforms3d==0.4.2 transparent-background==1.3.4