was this function copied from somewhere? if so, what changes were made?

it's a copy of nemo/tron/model.py:
https://github.com/NVIDIA/NeMo/blob/8ddf4387344c6423763ec9ee0c9a755cbb5d8d35/nemo/tron/model.py

the main change is removing the Float16Module wrapper (which is what originally casts the model logits output to float32):
https://github.com/NVIDIA-NeMo/RL/pull/856/files/a020289609cfa0d7a695a175eed009fdb4695088#diff-37539801eab6c58172c5cf85be33a1f9eac04c096a8e23170550ddf3bff8e3b3R125-R128

If it's only a one-line change, I'd prefer the change be reflected in the submodule (you can branch where the submodule is at to update)

also, if you expect the model coming back to not be a FP16 but something else, could you add a test asserting the model type? We're currently migrating away from tron, so once that is done, this test would ensure we don't miss this typing fix you're adding

updated NeMo submodule

branch: https://github.com/NVIDIA/NeMo/tree/pjin/nemorl-logprob
commit: NVIDIA-NeMo/NeMo@0bf0dbc

also, if you expect the model coming back to not be a FP16 but something else, could you add a test asserting the model type? We're currently migrating away from tron, so once that is done, this test would ensure we don't miss this typing fix you're adding

what I did is add a float32 dtype check to the existing megatron logprobs test, and running that test on more cases of (logprob chunk size, deferred float32 logits)
https://github.com/NVIDIA-NeMo/RL/pull/856/files#diff-9556cb57e37308923c54e7a6df8982afafef5e36544f350af3324db43f74bdbeR703

one thing is that the policy model worker mainly exposes the model output through get_logprobs, and there is not another interface for getting at the underlying torch model logits. but I think just checking that the returned logprobs are float32 should be sufficient?

can we add a unit test for this function so we make sure the non-chunked version equals the chunk?

added a chunk_size parameter to DistributedLogprobTestActor:

https://github.com/NVIDIA-NeMo/RL/pull/856/files#diff-13516af2fb4ffe4b66dd0cc8a1113bcd8334a1500f77ab0ce0f503bf529d8dc9

what would be the reason to set this to False?

mostly for strict backward compat, but we could instead enable it by default (i.e. make it an opt-out config like no_defer_fp32_logits or similar)

wdyt?

I see. How about the following:

1. this PR introduces it, default off
2. follow up PR where we run all our nightly tests to see if defaulting to true is ok, if so, remove the arg
   wdyt? If the feature is broadly applicable we should probably switch it to true so no one else runs into the same issue (assuming no accuracy penalty)

yup, (1) and then (2) SGTM!