Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| label = torch.tensor(batch_label) | ||
| hw_list = torch.tensor(hw_list, dtype=torch.int32) | ||
|
|
||
| # # Move tensors to model device for FSDP2 compatibility | ||
| # device = next(self.model.parameters()).device |
There was a problem hiding this comment.
During NitTrainer.compute_loss the labels and hw_list remain on CPU (torch.tensor(...) with no .to(device) after the move block was commented out), while the latents/noise and model live on CUDA. When FlowMatchingLoss calls the NiT model, the label embeddings and rotary cache use these CPU tensors, which will raise a device mismatch as soon as training runs on GPU. Please move label and hw_list to the model device before building model_kwargs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| label = torch.tensor(batch_label) | ||
| hw_list = torch.tensor(hw_list, dtype=torch.int32) | ||
|
|
There was a problem hiding this comment.
The new NiT trainer builds label and hw_list tensors on CPU but then passes them directly into self.loss_fn/self.model while images and latents are on CUDA. NitModel embeds y and computes RoPE grids inside the forward pass, so CPU indices against CUDA buffers will raise a device mismatch as soon as training starts. The commented-out .to(device) calls just below suggest these tensors were intended to be moved. Please transfer label and hw_list to the model device before invoking the model to avoid runtime crashes when running on GPUs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Is this part nvidia radio files all necessary? Would be nice if we can remove some parts if the modeling nit is not actually using it
…o prevent Hydra conflicts.
… calculate token dimensions and return class ID.
…tDataProcessor` to accept a dictionary row.
… token counts in NitDataset.
…tiple sequences into a single batch.
…s and update package dependencies.
…ating type hints in NitDataset and NitProcessor.
…, and modify example configuration for dataset format and output directory.
… packing strategy updates
… directory and logging settings
… eliminate unused modeling_utils module
…itTrainer for enhanced model configuration and loss computation
…in prepare_nit function
…ding gradient checkpointing and attention implementation adjustments
…ainer with additional parameters for improved training dynamics
3 participants
Currently support c2i