Skip to content

fix: preserve batch-major LAM action ordering - #25

Open
zjn20030811 wants to merge 1 commit into
NVIDIA:mainfrom
zjn20030811:fix/preserve-lam-batch-order
Open

fix: preserve batch-major LAM action ordering#25
zjn20030811 wants to merge 1 commit into
NVIDIA:mainfrom
zjn20030811:fix/preserve-lam-batch-order

Conversation

@zjn20030811

Copy link
Copy Markdown

Summary

  • restore LAM latent actions with the batch-major layout used by lam_video
  • add a dependency-light regression test using unique per-sample/action tokens

Motivation

forward flattens lam_video as (b p) before encoding. The subsequent inverse reshape interpreted that same flattened dimension as (t b), which silently interleaved action tokens across samples whenever both batch size and action length were greater than one. The tensor shape remained valid, so the issue was not visible in shape-only checks.

Validation

  • python -m pytest -q --confcutdir=tests tests/test_lam_batch_order.py (2 passed)
  • python -m py_compile tests/test_lam_batch_order.py cosmos_predict2/_src/predict2/models/text2world_model_rectified_flow.py
  • git diff --check

Fixes #23.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Inconsistent batch/time reshape may mix LAM actions across samples during pretraining

1 participant