[RLlib] Make sure SlateQ works with GPU. #22738
Conversation
…with GPU training.
| @@ -41,14 +41,15 @@ def build_slateq_model_and_distribution( | |||
| Returns: | |||
| Tuple consisting of 1) Q-model and 2) an action distribution class. | |||
| """ | |||
| device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") | |||
There was a problem hiding this comment.
Actually, the TorchPolicy will take care of all this.
I think we only have to make sure all tensors in the loss function are on the right device. ...
There was a problem hiding this comment.
ok, I understand everything now.
it's the target_model that was the issue. we just need to use the correct target_model out of policy.target_models for things to work.
thanks.
There was a problem hiding this comment.
Ah, yeah, sorry, I should have thought about this. Yes, you can always do:
correct_target_model_to_use = policy.target_models[model]
...
| @@ -154,7 +155,7 @@ def build_slateq_losses( | |||
|
|
|||
| clicked = torch.sum(click_indicator, dim=1) | |||
| mask_clicked_slates = clicked > 0 | |||
| clicked_indices = torch.arange(batch_size) | |||
| clicked_indices = torch.arange(batch_size).to(policy.device) | |||
There was a problem hiding this comment.
Here, this seems (almost) correct.
clicked_indices = torch.arange(batch_size).to(clicked.device). # <- some tensor that we know is already on one of the GPU.
There was a problem hiding this comment.
yep. this was actually how I did it originally. :)
updated.
| @@ -320,7 +321,10 @@ def score_documents( | |||
| torch.multiply(user_obs.unsqueeze(1), torch.stack(doc_obs, dim=1)), dim=2 | |||
| ) | |||
| # Compile a constant no-click score tensor. | |||
| score_no_click = torch.full(size=[user_obs.shape[0], 1], fill_value=no_click_score) | |||
| # Make sure it lives on the same device as scores_per_candidate. | |||
| score_no_click = torch.full( | |||
| # [1, AxS] Useful for torch.take_along_dim() | ||
| policy.slates_indices = policy.slates.reshape(-1).unsqueeze(0).to(policy.device) | ||
|
|
||
| setup_mixins(policy) |
There was a problem hiding this comment.
ah, reverted. I was trying to move policy.slates_indices to the correct device during late_setup.
but I am doing this the correct way now.
sven1977
changed the title
Why are these changes needed?
Create models and variables on proper device so SlateQ works with GPU training.
Related issue number
Checks
scripts/format.shto lint the changes in this PR.