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

[guijuzhejiang]

🐛 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