OOM while load the whole model

[Milesld]

Thanks for your great work.

I meet an OOM error while load the model. After check the code, it seems can be fixed in this file: https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/cosmos_predict2/_src/predict2/utils/model_loader.py#L366
Here .cuda() was applied first before apply_fsdp, becasuse of the big model, it cause OOM. Can I change the .cuda to end after apply_fsdp? I change the code and it seems work fine, but I have no idea is this ok...The code shows below:
`
def create_model_from_consolidated_checkpoint_with_fsdp(config: Config) -> ImaginaireModel:

"""
Instantiate a model, load weights from a consolidated checkpoint, and initialize FSDP if required.

Args:
    config: The configuration object for the experiment.

Returns:
    model: The loaded and (optionally) FSDP-wrapped model.
"""
# To avoid DTensor issues, load the model from a consolidated checkpoint in Tensor format before applying FSDP.
fsdp_shard_size = config.model.config.fsdp_shard_size
config.model.config.fsdp_shard_size = 1  # Set to 1 to disable FSDP during model instantiation.
# model = instantiate(config.model).cuda()
model = instantiate(config.model)
# DCP checkpointer does not support loading from a consolidated checkpoint, so we support it here.
model = load_model_state_dict_from_checkpoint(
    model=model,
    config=config,
    s3_checkpoint_dir=config.checkpoint.load_path,
    load_ema_to_reg=config.checkpoint.load_ema_to_reg,
)
# If FSDP is enabled, apply FSDP to the model.
if fsdp_shard_size > 1:
    config.model.config.fsdp_shard_size = fsdp_shard_size
    fsdp_device_mesh = hsdp_device_mesh(
        sharding_group_size=fsdp_shard_size,
    )
    if hasattr(model, "apply_fsdp") and callable(model.apply_fsdp):
        model.apply_fsdp(fsdp_device_mesh)
    else:
        raise AttributeError(
            "Model does not implement 'apply_fsdp'. Please implement this method to enable FSDP after consolidated checkpoint loading."
        )
model = model.cuda()

return model

`

And another question, is the loss value between 0~0.5 (I just trained for about 100 steps)? And there is no downward trend, just up and down between

[jcl2023]

@Milesld

I have rtx5070ti which has 16GB vram and got OOM when run inference of the bus terminal example,
Any idea how much vram is needed to run an inference?

[ccccc5]

@jcl2023
The default backend model is cosmos-reason1 7B,
Replace it with cosmos-reason2 2B and it will work.

[jcl2023]

@ccccc5

Do you have an instruction of how to replace cosmos-reason1 7B with cosmos-reason2 2B?

Thanks,

[ccccc5]

@jcl2023
Here's my repo with said modifications : https://github.com/ccccc5/cosmos-predict2.5-r2
It is an "unofficial wip" and is not yet producing good results (perhaps needs finetuning with reason2 2B) but the pipeline works happily with my 16GB card

Here's an example command (post you setup the cosmos requirements with uv):

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
python examples/inference.py \
  -i assets/base/bus_terminal.json \
  -o outputs/base_bus_terminal_i2w  \
  --model=2B/post-trained \
  --inference-type=image2world \
  --offload-diffusion-model \
  --offload-tokenizer \
  --offload-text-encoder \
  --disable-guardrails \
  --num-output-frames 15 \
  --num-steps 20 \
  --resolution 512,512

[jcl2023]

@ccccc5

Your repo has only license file.

[ccccc5]

@jcl2023
Something went wrong, this one should do https://github.com/ccccc5/cosmos-predict2.5#

Seems reason2-2B needs 24 GB to work as quoted here, so I'm not sure what/how is it working, will investigate in future

[jcl2023]

@ccccc5

There is no mention of what changes you do in this repo

[jcl2023]

@ccccc5
What changes have you done to use cosmos reason 2B in this repo?
https://github.com/ccccc5/cosmos-predict2.5#