Skip to content

Fix some issues in the code. #24

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ChiGod0502 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix some issues in the code. #24

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a3c4348
Refactor model checkpoint saving logic
ChiGod0502 Nov 22, 2025
b45cfda
Remove bos and eos tokens from hidden states when test
ChiGod0502 Nov 22, 2025
1bb4ef5
Remove bos and eos tokens from hidden states when training
ChiGod0502 Nov 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions prismatic/extern/hf/modeling_prismatic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -848,6 +848,7 @@ def _regression_or_discrete_prediction(
multi_layer_hidden_states = []

for item in language_model_output.hidden_states[0:]:
item = [:, 1:-1, :] # remove bos and eos token first
# last_hidden_states = output.hidden_states[-1] # (B, seq_len, D)
# Get hidden states for text portion of prompt+response (after the vision patches)
text_hidden_states = item
Expand Down
33 changes: 17 additions & 16 deletions vla-scripts/finetune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,9 +396,10 @@ def run_forward_pass(
multi_layer_hidden_states = []

for item in output.hidden_states[0:]:
item = item[:, 1:-1, :] # remove bos token and eos token
# last_hidden_states = output.hidden_states[-1] # (B, seq_len, D)
# Get hidden states for text portion of prompt+response (after the vision patches)
text_hidden_states = item[:, num_patches:-1]
text_hidden_states = item[:, num_patches:, :]
# Get hidden states for action portion of response
batch_size = batch["input_ids"].shape[0]
# actions_hidden_states = text_hidden_states[:, -1, :].reshape(batch_size, 1, -1).to(torch.bfloat16)
Expand Down Expand Up @@ -1081,21 +1082,21 @@ def rename_state_dict_keys(state_dict, replace_map):
optimizer.zero_grad()
progress.update()

# Save model checkpoint: either keep latest checkpoint only or all checkpoints
if gradient_step_idx > 0 and log_step % cfg.save_freq == 0:
save_training_checkpoint(
cfg=cfg,
run_dir=run_dir,
log_step=log_step,
vla=vla,
processor=processor,
proprio_projector=proprio_projector if cfg.use_proprio else None,
noisy_action_projector=None,
action_head=action_head,
train_dataset=train_dataset,
distributed_state=distributed_state,
new_state_dict=RAW_STATE_DICT,
)
# Save model checkpoint: either keep latest checkpoint only or all checkpoints
if gradient_step_idx > 0 and log_step % cfg.save_freq == 0:
save_training_checkpoint(
cfg=cfg,
run_dir=run_dir,
log_step=log_step,
vla=vla,
processor=processor,
proprio_projector=proprio_projector if cfg.use_proprio else None,
noisy_action_projector=None,
action_head=action_head,
train_dataset=train_dataset,
distributed_state=distributed_state,
new_state_dict=RAW_STATE_DICT,
)

# Test model on validation set
if cfg.use_val_set and log_step > 0 and log_step % cfg.val_freq == 0:
Expand Down