Um, why do we need cross entropy in get_per_token_logps?

Apparently, these will return logprobs and are equivalent to selective softmax? But I am not sure if we want to return logprobs in this matrix because like @danielhanchen said we folded it into a torch.compile kernel. I am questioning if the memory saved here is actually from the cut cross entropy loss rather than the chunked concatenation of the hidden states. I am currently at work but we can check later if chunking the hidden states conserves similar amounts of memory.

Another thing I see is that according to to this person's post, there also seems to be some speed up as well, what we can do instead of materializing the logits outside of here is also put the linear_cut_cross_entropy in place of the code in selective_softmax so we get speed up and memory and do not materialize logits outside of the kernel.

@pluesclues That would work too. The only reason I did it this way is to ensure consistency with HF. That being said, we may, at some point, need to write a custom kernel anyway to run fused operations on the logit matrix chunk. Currently, the implementation in HF scales the logits with temperature before computing logps (https://github.com/huggingface/trl/blob/4c92de00001379ceedaf073512ce4df5da304d08/trl/trainer/grpo_trainer.py#L871).

Okay I just tested this method inside of the kernel, its as I suspected, we cannot use linear_cross_entropy which is a torch.compile kernel in itself inside of a torch.compile kernel, I confirmed this by running ref = -1 * linear_cross_entropy(ref_hidden_states_j[:, :-1, :].to(dtype=lm_head.dtype), lm_head, input_ids_j, reduction="none", impl="cce") right before accumulate_chunk outside the kernel and also called this line inside the kernel, outside the kernel it works just fine, inside it seems to break. I still haven't tested the speed up on my machine yet, but so far it looks like we can either merge this or just change the way we calculate logprobs to exactly how CCE does it in their kernel.

About the memory-saving speed-up I reported, I believe the manner in which I profiled it does not provide an accurate account. I am only logging the peak memory allocated throughout a training step, clearing it at the beginning. This approach fails to account for the memory allocated for the old and ref policies as they are computed and cached outside the new policy update loop, i.e., every _step % num_interations == 0. I expected much higher memory savings. I would appreciate some help with this.

Moreover, I would like to confirm that UNSLOTH_USE_NEW_MODEL being set to 0 must be interpreted as the pathway to UnslothEfficientGRPO as is the case in the current implementation.
Also, UNSLOTH_RETURN_HIDDEN_STATES is set to 1 before executing the forward pass in _get_per_token_logps but never reset to its original value, creating an unintended side-effect. This is done in a couple of places. Would it not be better to reset it?

Do you have the wandb of memory usage over time (as tracked by trl/wandb itself) of the run?

This is a much smaller run with double the batch size. The CCE version completes its 4 training steps in 7 mins, whereas the current implementation OOMs on my machine after 12 mins.
In this case, the amount of memory saved is roughly 25%.

batch_size = 16
unsloth_num_chunks = 4