Skip to content

Commit

Permalink
[TT-Train] fp32 turned off for softmax (#17683)
Browse files Browse the repository at this point in the history
### Problem description
fp32_dest doesn't work for softmax and log_softmax.
Training is exloding.

### What's changed
Disabled it for both ops.

### Checklist
- [x] [All post
commit](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml)
CI passes
https://github.com/tenstorrent/tt-metal/actions/runs/13189909685
- [x] New/Existing tests provide coverage for changes
  • Loading branch information
dmakoviichuk-tt authored Feb 7, 2025
1 parent e3920bd commit e7e86d7
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion tt-train/sources/ttml/core/compute_kernel_config.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ ttnn::WormholeComputeKernelConfig ComputeKernelConfig::precise() {

ttnn::WormholeComputeKernelConfig ComputeKernelConfig::softmax() {
ttnn::WormholeComputeKernelConfig config;
config.fp32_dest_acc_en = true;
config.fp32_dest_acc_en = false;
config.math_approx_mode = false;
config.math_fidelity = MathFidelity::HiFi4;
config.packer_l1_acc = true;
Expand Down
2 changes: 1 addition & 1 deletion tt-train/sources/ttml/ops/unary_ops.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ autograd::TensorPtr log_softmax_moreh(const autograd::TensorPtr& tensor, int dim
ttnn::operations::moreh::moreh_softmax::MorehSoftmaxOp::LOGSOFTMAX,
ttnn::operations::moreh::moreh_softmax::MorehSoftmaxOpParallelizationStrategy::NONE,
/* output_mem_config */ std::nullopt,
/* compute_kernel_config */ core::ComputeKernelConfig::precise());
/* compute_kernel_config */ core::ComputeKernelConfig::softmax());
auto out = autograd::create_tensor(log_softmax);

autograd::GradFunction grad = [tensor, out, dim]() {
Expand Down

0 comments on commit e7e86d7

Please sign in to comment.