Skip to content

Bump 3rdparty/NeMo from b685967 to 7d7a10c #807

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dependabot wants to merge 3 commits into from
Closed

Bump 3rdparty/NeMo from b685967 to 7d7a10c #807

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
607b19e
Bump 3rdparty/NeMo from `b685967` to `7d7a10c`
dependabot[bot] Apr 7, 2025
3ac72f6
remove call to context_parallel loss
pstjohn Apr 7, 2025
aad6987
bump NeMo to TOT
pstjohn Apr 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rdparty/NeMo
Open in desktop
Submodule NeMo updated from b68596 to 7d7a10
17 changes: 2 additions & 15 deletions sub-packages/bionemo-geneformer/src/bionemo/geneformer/model/finetune_token_regressor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,7 @@
from nemo.collections.llm.peft.lora import LoRA, LoRALinear
from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import ParallelLinearAdapter
from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group
from nemo.lightning.megatron_parallel import (
masked_token_loss,
masked_token_loss_context_parallel,
)
from nemo.lightning.megatron_parallel import masked_token_loss
from torch import Tensor, nn

from bionemo.llm.model.biobert.model import BioBertConfig, BioBertOutput, MegatronBioBertModel
Expand Down Expand Up @@ -102,17 +99,7 @@ def forward(
# TODO(@jstjohn) also handle different output keys, like the sequence loss.

cp_size = parallel_state.get_context_parallel_world_size()
if cp_size == 1:
# reduce the loss across the micro batch
loss_for_microbatch = masked_token_loss(unreduced_token_loss, batch["loss_mask"])
else:
# reduce the loss across the micro batch.
# TODO(@jomitchell): Figure out who defines "num_valid_tokens_in_ub" in the batch and document/understand this.
# This has something to do with context parallel, and there is probably a megatron or nemo function that adds this and
# other necessary keys to the batch. Thanks!
loss_for_microbatch = masked_token_loss_context_parallel(
unreduced_token_loss, batch["loss_mask"], batch["num_valid_tokens_in_ub"]
)
loss_for_microbatch = masked_token_loss(unreduced_token_loss, batch["loss_mask"], cp_size)

# If we do not drop the last partial batch of validation, we need to do fancy reduction handling to support
# reducing the loss across the data parallel group.
Expand Down
19 changes: 3 additions & 16 deletions sub-packages/bionemo-llm/src/bionemo/llm/model/loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,7 @@
from megatron.core import parallel_state, tensor_parallel
from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group
from nemo.lightning.megatron_parallel import (
MegatronLossReduction,
masked_token_loss,
masked_token_loss_context_parallel,
)
from nemo.lightning.megatron_parallel import MegatronLossReduction, masked_token_loss
from torch import Tensor


Expand Down Expand Up @@ -181,17 +177,8 @@ def forward(

# compute loss
cp_size = parallel_state.get_context_parallel_world_size()
if cp_size == 1:
# reduce the loss across the micro batch per valid token
loss_for_microbatch = masked_token_loss(unreduced_token_loss, batch["loss_mask"])
else:
# reduce the loss across the micro batch per valid token.
# TODO(@jomitchell): Figure out who defines "num_valid_tokens_in_ub" in the batch and document/understand this.
# This has something to do with context parallel, and there is probably a megatron or nemo function that adds this and
# other necessary keys to the batch. Thanks!
loss_for_microbatch = masked_token_loss_context_parallel(
unreduced_token_loss, batch["loss_mask"], batch["num_valid_tokens_in_ub"]
)
# reduce the loss across the micro batch per valid token
loss_for_microbatch = masked_token_loss(unreduced_token_loss, batch["loss_mask"], cp_size)

# If we do not drop the last partial batch of validation, we need to do fancy reduction handling to support
# reducing the loss across the data parallel group.
Expand Down