Fix torch.compile crash in generating triton code for normalise_grad

[lkhphuc]

dtype mismatch between tensor and python float when pytorch generate this triton code.

  triton.compiler.errors.CompilationError: at 7:11:
  def triton_poi_fused_mul_0(in_ptr0, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
      xoffset = tl.program_id(0) * XBLOCK
      xindex = xoffset + tl.arange(0, XBLOCK)[:]
      xmask = xindex < xnumel
      x0 = xindex
      tmp0 = tl.load(in_ptr0 + (x0), xmask).to(tl.float32)
      tmp1 = libdevice.pow(ks0 / ks1, tl.full([], 0.500000000000000, tl.float64))
             ^
  (triton.language.float32 ,triton.language.float64)

Fixed it by explicitly convert to Tensor and uses torch primitive.
I'm using Triton comes with Pytorch nightly btw.

[rakkit]

@lkhphuc Thanks a lot and sry for some reason I did not see this PR.

So the problem appears to be the full_graph compile decorator here https://github.com/SDLAML/disco/blob/main/disco/abstract_disco.py#L92

After some point in the latest PyTorch version, it is broken.

[rakkit]

Thanks a lot. sry again for the late merge. I did test and it works well for me. (torch 2.10.0.dev20251124+cu129)

[rakkit]

hey @lkhphuc

i think i have to revert these changes, not your code's problem, but compile. (i will find time after vacation to refactor code to fix it )

There is a stupid and annoying bug in Torch compile.

The current logic for LLM's embedding/head is, AT most of the time, we do lmo on a sharded tensor. (because lmo for head/output is row-wise for [vocab, dim] parameters, such that we dont have to gather it)

BUT, when we try to "track" the norm of the update/weights of head/output, we gather it and run lmo on full tensor. Therefore, at the very early of beginning, most of the lmo calls receive the sharded tensor of embedding with a few times a full tensor.

This works well until a certain step, about 260-270 steps, the compile decides to replace the lmo kernel for norm track with lmo kernel for gradient update (the constant variable of matrix shape reduced by world size) that the norm tracker will break at step 260~270.

(good thing is it seems only affect to logging for tracking, no bother with actual training dynamic)