initial grok

[dan-garvey]

Initial grok work, also does some refactoring

[archana-ramalingam]

Posted a comment here, doesn't show up on the main page.

[dan-garvey]

Posted a comment here, doesn't show up on the main page.

lets chat at meeting

[dan-garvey]

I did say I didn't mind if we did this, but looking at it now, all the args are just values from the config anyway. I feel like separating them I think just makes our code harder to follow.

Not a big deal either way, but I'd personally prefer we just drop this commit. @KyleHerndon what do you think?

[KyleHerndon]

I think this is the right way to set things up. The bad argument is that llama.cpp has a similar setup, and we're referencing them for accuracy baselines. The better argument goes something like:

The KV cache is not ontologically related to the config, it just (currently) exclusively uses parameters from it. In the near future, we will want things like sharding, at which point the KV Cache might need additional args (number of devices, for example, which is not a ModelConfig parameter, but something like an ExecutionConfig parameter).

This change doesn't feel related to grok as much a code refactor, which I think is fine to include in a bigger patch for efficiency but might make things harder to follow when looking at commit history. Otherwise, I find the code just as easy to follow, if not slightly easier because I think it is better organized.

[archana-ramalingam]

I don't mind either way, but do agree with Kyle that this keeps it better organized. We can consult Rob or Stella to see if there is a better way to do this.

[dan-garvey]

not necessary, if you two agree thats more than enough for me

[dan-garvey]

looks good to me, @KyleHerndon @archana-ramalingam any final changes you think are needed?

[dan-garvey]

nice catch

[archana-ramalingam]

I tested Llama and Grok models again and it's LGTM.

[sogartar]

@dan-garvey, the error.
Seems to be related to calling __mul__ on the DefaultPrimitiveTensor.
This fails:

output = weight * to(output, weight.dtype)

but this is fine

output = elementwise(torch.mul, weight, to(output, weight.dtype))

__mul__'s implementation is just

    def __mul__(self, rhs):
        from ..ops import elementwise

        return elementwise(torch.mul, self, rhs)

[sogartar]

When calling torch.export.export with strict=False, the error is not present, so it is about the export checks.

[sogartar]

Instead of calling the * operator I changed the code to call a custom member function in weight

output = weight.my_mul(to(output, weight.dtype))

class InferenceTensor(ABC):

....

    def my_mul(self, rhs):
        from ..ops import elementwise

        return elementwise(torch.mul, self, rhs)

This is essentially the same code, but PyTorch does something special about binary operators. '+' also suffers from the same problem.

Also when running with env vars

TORCH_LOGS="+dynamo"
TORCHDYNAMO_VERBOSE=1

The tracer reports

V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:798] [0/0] [__trace_bytecode] TRACE LOAD_FAST results []
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:798] [0/0] [__trace_bytecode] TRACE LOAD_CONST 0 [TupleVariable()]
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:798] [0/0] [__trace_bytecode] TRACE BINARY_SUBSCR None [TupleVariable(), ConstantVariable()]
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:798] [0/0] [__trace_bytecode] TRACE RETURN_VALUE None [TensorVariable()]
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:2807] [0/0] DONE INLINING <code object __call__ at 0x7f618b39f470, file "/home/bpetkant/ws/sharktank/repo/sharktank/sharktank/ops/_registry.py", line 196>
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:798] [0/0] [__trace_bytecode] TRACE BINARY_MULTIPLY None [UserDefinedObjectVariable(), TensorVariable()]
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:814] [0/0] empty checkpoint
V0926 06:19:31.456000 140060607524864 torch/_dynamo/symbolic_convert.py:2796] [0/0] FAILED INLINING <code object rms_norm_default at 0x7f618b2126b0, file "/home/bpetkant/ws/sharktank/repo/sharktank/sharktank/ops/default_impls.py", line 308>

It is doing something special about BINARY_MULTIPLY.

[sogartar]

I opened an issue with PyTorch. I think that may be a bug there.